This pull request introduces several enhancements and database schema changes to support associating LLM completions with prospects, improve API security, and add utility scripts for database management.

.env.sample

@@ -1,3 +1,4 @@
+PYTHON_KEY=
 GEMINI_API_KEY=
 RESEND_API_KEY=
 BASE_URL=

app/api/llm/llm.py

@@ -1,7 +1,8 @@
 import os
-from fastapi import APIRouter, HTTPException, Query, Request
+from fastapi import APIRouter, HTTPException, Query, Request, Depends
 from app.utils.make_meta import make_meta
 from app.utils.db import get_db_connection_direct
+from app.utils.api_key_auth import get_api_key
 
 router = APIRouter()
 
@@ -10,6 +11,7 @@ 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")
+    , api_key: str = Depends(get_api_key)
 ) -> dict:
     """GET /llm: Paginated list of LLM completions."""
     try:
@@ -20,7 +22,7 @@ def get_llm_records(
         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
+            SELECT id, prompt, completion, duration, time, data, model, prospect_id
             FROM llm
             ORDER BY id DESC
             LIMIT %s OFFSET %s;
@@ -34,6 +36,7 @@ def get_llm_records(
                 "time": row[4].isoformat() if row[4] else None,
                 "data": row[5],
                 "model": row[6],
+                "prospect_id": row[7],
             }
             for row in cur.fetchall()
         ]
@@ -58,6 +61,7 @@ def get_llm_records(
 def llm_post(payload: dict) -> dict:
     """POST /llm: send prompt to Gemini, returns completion google-genai SDK."""
     prompt = payload.get("prompt")
+    prospect_id = payload.get("prospect_id")
     if not prompt:
         raise HTTPException(status_code=400, detail="Missing 'prompt' in request body.")
     api_key = os.getenv("GEMINI_API_KEY")
@@ -105,11 +109,11 @@ def llm_post(payload: dict) -> dict:
             cur = conn.cursor()
             cur.execute(
                 """
-                INSERT INTO llm (prompt, completion, duration, data, model)
-                VALUES (%s, %s, %s, %s, %s)
+                INSERT INTO llm (prompt, completion, duration, data, model, prospect_id)
+                VALUES (%s, %s, %s, %s, %s, %s)
                 RETURNING id;
                 """,
-                (prompt, completion, duration, data_blob, used_model)
+                (prompt, completion, duration, data_blob, used_model, prospect_id)
             )
             record_id_row = cur.fetchone()
             record_id = record_id_row[0] if record_id_row else None

app/api/llm/sql/alter_add_prompt_code.sql

@@ -0,0 +1,2 @@
+-- Migration: Add prompt_code column to llm table
+ALTER TABLE llm ADD COLUMN IF NOT EXISTS prompt_code TEXT;

app/api/llm/sql/alter_add_prospect_id.sql

@@ -0,0 +1,2 @@
+-- Migration: Add prospect_id column to llm table
+ALTER TABLE llm ADD COLUMN IF NOT EXISTS prospect_id INTEGER REFERENCES prospects(id);

app/api/llm/sql/create_table.sql

@@ -7,5 +7,6 @@ CREATE TABLE IF NOT EXISTS llm (
 	duration FLOAT,
 	time TIMESTAMPTZ DEFAULT NOW(),
 	data JSONB,
-	model TEXT
+	model TEXT,
+	prospect_id INTEGER REFERENCES prospects(id)
 );

app/api/llm/sql/empty_llm_table.py

@@ -0,0 +1,14 @@
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../..')))
+from app.utils.db import get_db_connection_direct
+
+if __name__ == "__main__":
+    sql = "DELETE FROM llm;"
+    conn = get_db_connection_direct()
+    cur = conn.cursor()
+    cur.execute(sql)
+    conn.commit()
+    cur.close()
+    conn.close()
+    print("All records deleted from llm table.")

app/api/llm/sql/run_alter_add_prompt_code.py

@@ -0,0 +1,16 @@
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../..')))
+from app.utils.db import get_db_connection_direct
+
+if __name__ == "__main__":
+    sql = """
+    ALTER TABLE llm ADD COLUMN IF NOT EXISTS prompt_code TEXT;
+    """
+    conn = get_db_connection_direct()
+    cur = conn.cursor()
+    cur.execute(sql)
+    conn.commit()
+    cur.close()
+    conn.close()
+    print("Migration complete: prompt_code column added to llm table.")

app/api/llm/sql/run_alter_add_prospect_id.py

@@ -0,0 +1,16 @@
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../..')))
+from app.utils.db import get_db_connection_direct
+
+if __name__ == "__main__":
+    sql = """
+    ALTER TABLE llm ADD COLUMN IF NOT EXISTS prospect_id INTEGER REFERENCES prospects(id);
+    """
+    conn = get_db_connection_direct()
+    cur = conn.cursor()
+    cur.execute(sql)
+    conn.commit()
+    cur.close()
+    conn.close()
+    print("Migration complete: prospect_id column added to llm table.")

app/api/prospects/prospects.py

@@ -184,6 +184,29 @@ def prospects_read_one(id: int = Path(..., description="ID of the prospect to re
             if row is not None:
                 columns = [desc[0] for desc in cur.description]
                 data = dict(zip(columns, row))
+                # Fetch related llm records
+                try:
+                    from app.utils.db import get_db_connection_direct
+                    llm_conn = get_db_connection_direct()
+                    llm_cur = llm_conn.cursor()
+                    llm_cur.execute("SELECT id, prompt, completion, duration, time, data, model FROM llm WHERE prospect_id = %s ORDER BY id DESC;", (id,))
+                    llm_records = [
+                        {
+                            "id": r[0],
+                            "prompt": r[1],
+                            "completion": r[2],
+                            "duration": r[3],
+                            "time": r[4].isoformat() if r[4] else None,
+                            "data": r[5],
+                            "model": r[6],
+                        }
+                        for r in llm_cur.fetchall()
+                    ]
+                    llm_cur.close()
+                    llm_conn.close()
+                    data["llm_records"] = llm_records
+                except Exception as llm_exc:
+                    data["llm_records"] = []
             else:
                 data = None
                 meta = make_meta("error", f"No prospect found with id {id}")

app/api/prospects/sql/create_table.sql

@@ -0,0 +1,40 @@
+CREATE TABLE IF NOT EXISTS prospects (
+	id SERIAL PRIMARY KEY,
+	first_name TEXT,
+	last_name TEXT,
+	title TEXT,
+	company_name TEXT,
+	email TEXT,
+	email_status TEXT,
+	primary_email_source TEXT,
+	primary_email_verification_source TEXT,
+	email_confidence TEXT,
+	primary_email_catchall_status TEXT,
+	primary_email_last_verified_at TEXT,
+	seniority TEXT,
+	sub_departments TEXT,
+	work_direct_phone TEXT,
+	home_phone TEXT,
+	mobile_phone TEXT,
+	corporate_phone TEXT,
+	other_phone TEXT,
+	do_not_call TEXT,
+	lists TEXT,
+	person_linkedin_url TEXT,
+	country TEXT,
+	subsidiary_of TEXT,
+	subsidiary_of_organization_id TEXT,
+	tertiary_email TEXT,
+	tertiary_email_source TEXT,
+	tertiary_email_status TEXT,
+	tertiary_email_verification_source TEXT,
+	primary_intent_topic TEXT,
+	primary_intent_score TEXT,
+	secondary_intent_topic TEXT,
+	secondary_intent_score TEXT,
+	qualify_contact TEXT,
+	cleaned TEXT,
+	search_vector TSVECTOR,
+	flag BOOLEAN DEFAULT FALSE,
+	hide BOOLEAN DEFAULT FALSE
+);

app/api/prospects/sql/describe_prospects_table.py

@@ -0,0 +1,20 @@
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../..')))
+from app.utils.db import get_db_connection_direct
+
+if __name__ == "__main__":
+    sql = """
+    SELECT column_name, data_type, is_nullable, column_default
+    FROM information_schema.columns
+    WHERE table_name = 'prospects'
+    ORDER BY ordinal_position;
+    """
+    conn = get_db_connection_direct()
+    cur = conn.cursor()
+    cur.execute(sql)
+    columns = cur.fetchall()
+    cur.close()
+    conn.close()
+    for col in columns:
+        print(col)

app/api/root.py

@@ -20,11 +20,10 @@ def root() -> dict:
         "time": epoch,
     }
     endpoints = [
-        {"name": "llm", "url": f"{base_url}/llm"},
         {"name": "docs", "url": f"{base_url}/docs"},
         {"name": "resend", "url": f"{base_url}/resend"},
         {"name": "health", "url": f"{base_url}/health"},
-        {"name": "prompts", "url": f"{base_url}/prompts"},
         {"name": "prospects", "url": f"{base_url}/prospects"},
+        {"name": "llm", "url": f"{base_url}/llm"},
     ]
     return {"meta": meta, "data": endpoints}

app/utils/api_key_auth.py

@@ -0,0 +1,22 @@
+import os
+from fastapi import Request, HTTPException, status
+from app.utils.make_meta import make_meta
+from fastapi.security.api_key import APIKeyHeader
+from typing import Optional
+from dotenv import load_dotenv
+
+load_dotenv()
+
+API_KEY = os.getenv("PYTHON_KEY")
+API_KEY_NAME = "X-API-Key"
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
+
+def get_api_key(request: Request, api_key_header: Optional[str] = ""):
+    if API_KEY is None:
+        meta = make_meta("error", "PYTHON_KEY not set")
+        return {"meta": meta, "data": []}
+    if api_key_header == API_KEY:
+        return api_key_header
+    else:
+        meta = make_meta("error", "Invalid PYTHON_KEY Key")
+        return {"meta": meta, "data": []}