From e0758b2aa6b44043ef1e6f4f5d1432ec74a5cb80 Mon Sep 17 00:00:00 2001 From: Wei Zang Date: Sat, 4 Apr 2026 22:06:47 +0100 Subject: [PATCH 1/2] Add LLM DB schema, seed script & rename tests Create a new llm table SQL schema (vector(1536), prompt, completion, duration, time, data). Add a small psycopg2 seed script (insert_llm_lorem.py) that loads .env, connects to the DB and inserts 5 dummy records with random 1536-d vectors. Rename tests/test_gemini.py to tests/test_llm.py and update tests to use LLM naming and /llm endpoints, adjust skip/failure messages and assertions accordingly (including updated mocked API handling). These changes introduce the DB schema and local seeding helper and align tests to the LLM endpoint/name changes. --- app/api/llm/sql/create_table.sql | 10 ++++++++++ app/api/llm/sql/insert_llm_lorem.py | 25 +++++++++++++++++++++++++ tests/{test_gemini.py => test_llm.py} | 19 ++++++++++--------- 3 files changed, 45 insertions(+), 9 deletions(-) create mode 100644 app/api/llm/sql/create_table.sql create mode 100644 app/api/llm/sql/insert_llm_lorem.py rename tests/{test_gemini.py => test_llm.py} (78%) diff --git a/app/api/llm/sql/create_table.sql b/app/api/llm/sql/create_table.sql new file mode 100644 index 0000000..71fc624 --- /dev/null +++ b/app/api/llm/sql/create_table.sql @@ -0,0 +1,10 @@ + +CREATE TABLE IF NOT EXISTS llm ( + id SERIAL PRIMARY KEY, + vector vector(1536), + prompt TEXT NOT NULL, + completion TEXT NOT NULL, + duration FLOAT, + time TIMESTAMPTZ DEFAULT NOW(), + data JSONB +); diff --git a/app/api/llm/sql/insert_llm_lorem.py b/app/api/llm/sql/insert_llm_lorem.py new file mode 100644 index 0000000..ab0e9b8 --- /dev/null +++ b/app/api/llm/sql/insert_llm_lorem.py @@ -0,0 +1,25 @@ +import os +import psycopg2 +from dotenv import load_dotenv +import random + +load_dotenv() + +conn = psycopg2.connect( + host=os.getenv('DB_HOST'), + port=os.getenv('DB_PORT', '5432'), + dbname=os.getenv('DB_NAME'), + user=os.getenv('DB_USER'), + password=os.getenv('DB_PASSWORD'), +) +cur = conn.cursor() +lorem = 'Lorem ipsum dolor sit amet, consectetur adipiscing elit.' +for i in range(5): + cur.execute( + 'INSERT INTO llm (vector, prompt, completion, duration, data) VALUES (%s, %s, %s, %s, %s);', + ([random.random() for _ in range(1536)], lorem, lorem, random.uniform(0.1, 2.0), '{}') + ) +conn.commit() +print('Inserted 5 records.') +cur.close() +conn.close() diff --git a/tests/test_gemini.py b/tests/test_llm.py similarity index 78% rename from tests/test_gemini.py rename to tests/test_llm.py index 51f8e6d..d032dbf 100644 --- a/tests/test_gemini.py +++ b/tests/test_llm.py @@ -1,9 +1,9 @@ import os import pytest -def test_gemini_real_api(): +def test_llm_real_api(): api_key = os.getenv("GEMINI_API_KEY") if not api_key: - pytest.skip("GEMINI_API_KEY not set; skipping real Gemini API test.") + pytest.skip("GEMINI_API_KEY not set; skipping real LLM API test.") from google import genai client = genai.Client(api_key=api_key) try: @@ -14,7 +14,7 @@ def test_gemini_real_api(): completion = getattr(response, "text", None) assert completion is not None and "hello" in completion.lower() except Exception as e: - pytest.fail(f"Gemini real API call failed: {e}") + pytest.fail(f"LLM real API call failed: {e}") import sys import os sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))) @@ -26,16 +26,17 @@ def test_gemini_real_api(): -def test_gemini_get_endpoint(): - response = client.get("/gemini") + +def test_llm_get_endpoint(): + response = client.get("/llm") assert response.status_code == 200 data = response.json() assert "meta" in data assert data["meta"]["severity"] == "success" - assert "Gemini endpoint says hello" in data["meta"]["title"] + assert "LLM endpoint says hello" in data["meta"]["title"] -def test_gemini_post_endpoint(monkeypatch): +def test_llm_post_endpoint(monkeypatch): # Mock google-genai SDK to avoid real API call class MockGenAIResponse: text = "Test completion" @@ -50,12 +51,12 @@ class MockGenAIClient: monkeypatch.setattr("google.genai.Client", lambda *args, **kwargs: MockGenAIClient()) payload = {"prompt": "Test prompt"} - response = client.post("/gemini", json=payload) + response = client.post("/llm", json=payload) assert response.status_code == 200 data = response.json() assert "meta" in data assert data["meta"]["severity"] == "success" - assert "Gemini completion received" in data["meta"]["title"] + assert "completion received" in data["meta"]["title"] assert data["data"]["prompt"] == "Test prompt" assert data["data"]["completion"] == "Test completion" assert "data" in data From a49d6de84ac390a48147c7144a4f2e7e406e16ed Mon Sep 17 00:00:00 2001 From: Wei Zang Date: Sat, 4 Apr 2026 22:34:35 +0100 Subject: [PATCH 2/2] Add LLM pagination, DB logging & model column Add a paginated GET /llm endpoint that queries the llm table and returns page metadata. Enhance POST /llm to time Gemini calls, persist completions (prompt, completion, duration, data, model) to the database (errors logged but not exposed), and use a direct DB connection helper. Update SQL to add a model column to the llm table, bump package version to 2.1.2, and adjust tests to match the new meta title content. --- app/__init__.py | 2 +- app/api/llm/llm.py | 82 ++++++++++++++++++++++++++++---- app/api/llm/sql/create_table.sql | 3 +- tests/test_llm.py | 3 +- 4 files changed, 79 insertions(+), 11 deletions(-) diff --git a/app/__init__.py b/app/__init__.py index aa0c61d..6cc2148 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -1,4 +1,4 @@ """Python - FastAPI, Postgres, tsvector""" # Current Version -__version__ = "2.1.1" +__version__ = "2.1.2" diff --git a/app/api/llm/llm.py b/app/api/llm/llm.py index 7481a98..e8b08a2 100644 --- a/app/api/llm/llm.py +++ b/app/api/llm/llm.py @@ -1,14 +1,58 @@ import os -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, Query, Request from app.utils.make_meta import make_meta +from app.utils.db import get_db_connection_direct router = APIRouter() @router.get("/llm") -def root() -> dict: - """GET /llm endpoint.""" - meta = make_meta("success", "LLM endpoint says hello") - return {"meta": meta} +def get_llm_records( + request: Request, + page: int = Query(1, ge=1, description="Page number (1-based)"), + page_size: int = Query(10, ge=1, le=100, description="Records per page") +) -> dict: + """GET /llm: Paginated list of LLM completions.""" + try: + conn = get_db_connection_direct() + cur = conn.cursor() + offset = (page - 1) * page_size + cur.execute("SELECT COUNT(*) FROM llm;") + count_row = cur.fetchone() + total = count_row[0] if count_row and count_row[0] is not None else 0 + cur.execute(""" + SELECT id, prompt, completion, duration, time, data, model + FROM llm + ORDER BY id DESC + LIMIT %s OFFSET %s; + """, (page_size, offset)) + records = [ + { + "id": row[0], + "prompt": row[1], + "completion": row[2], + "duration": row[3], + "time": row[4].isoformat() if row[4] else None, + "data": row[5], + "model": row[6], + } + for row in cur.fetchall() + ] + cur.close() + conn.close() + meta = make_meta("success", f"LLM {len(records)} records (page {page})") + return { + "meta": meta, + "data": { + "page": page, + "page_size": page_size, + "total": total, + "pages": (total + page_size - 1) // page_size, + "data": records, + }, + } + except Exception as e: + meta = make_meta("error", f"DB error: {str(e)}") + return {"meta": meta, "data": {}} @router.post("/llm") def llm_post(payload: dict) -> dict: @@ -22,8 +66,8 @@ def llm_post(payload: dict) -> dict: import logging try: from google import genai + import time as time_mod client = genai.Client(api_key=api_key) - # Try a list of known Gemini model names model_names = [ "models/gemini-flash-latest", "models/gemini-1.5-pro", @@ -36,6 +80,7 @@ def llm_post(payload: dict) -> dict: completion = None used_model = None errors = {} + start_time = time_mod.time() for model_name in model_names: try: response = client.models.generate_content(model=model_name, contents=prompt) @@ -46,12 +91,33 @@ def llm_post(payload: dict) -> dict: except Exception as e: errors[model_name] = str(e) continue + duration = time_mod.time() - start_time if not completion: error_details = " | ".join([f"{k}: {v}" for k, v in errors.items()]) raise Exception(f"No available Gemini model succeeded for generate_content with your API key. Details: {error_details}") + # Insert record into llm table + try: + import json + from app import __version__ + data_blob = json.dumps({"version": __version__}) + conn = get_db_connection_direct() + cur = conn.cursor() + cur.execute( + """ + INSERT INTO llm (prompt, completion, duration, data, model) + VALUES (%s, %s, %s, %s, %s); + """, + (prompt, completion, duration, data_blob, used_model) + ) + conn.commit() + cur.close() + conn.close() + except Exception as db_exc: + # Log DB error but do not fail the API response + logging.error(f"Failed to insert llm record: {db_exc}") + meta = make_meta("success", f"Gemini completion received from {used_model}") + return {"meta": meta, "data": {"prompt": prompt, "completion": completion}} except Exception as e: meta = make_meta("error", f"Gemini API error: {str(e)}") return {"meta": meta, "data": {}} - meta = make_meta("success", f"Gemini completion received from {used_model}") - return {"meta": meta, "data": {"prompt": prompt, "completion": completion}} diff --git a/app/api/llm/sql/create_table.sql b/app/api/llm/sql/create_table.sql index 71fc624..41c47ec 100644 --- a/app/api/llm/sql/create_table.sql +++ b/app/api/llm/sql/create_table.sql @@ -6,5 +6,6 @@ CREATE TABLE IF NOT EXISTS llm ( completion TEXT NOT NULL, duration FLOAT, time TIMESTAMPTZ DEFAULT NOW(), - data JSONB + data JSONB, + model TEXT ); diff --git a/tests/test_llm.py b/tests/test_llm.py index d032dbf..b31fc8a 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -33,7 +33,8 @@ def test_llm_get_endpoint(): data = response.json() assert "meta" in data assert data["meta"]["severity"] == "success" - assert "LLM endpoint says hello" in data["meta"]["title"] + assert "LLM" in data["meta"]["title"] + assert "records" in data["meta"]["title"] def test_llm_post_endpoint(monkeypatch):