main.py

import os
import sqlite3
import json
import faiss
import numpy as np
import dashscope
from http import HTTPStatus
from fastapi import FastAPI, Request, HTTPException, status
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional
from openai import OpenAI
from dotenv import load_dotenv

# --- CONFIGURATION ---
load_dotenv()
DATA_DIR = "data"

# --- API Clients ---
client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)

# --- Dify API Models ---
class RetrievalSetting(BaseModel):
    top_k: int
    score_threshold: float

class RetrievalRequest(BaseModel):
    knowledge_id: str
    query: str
    retrieval_setting: RetrievalSetting
    metadata_condition: Optional[Dict] = None

class Record(BaseModel):
    content: str
    score: float
    title: str
    metadata: Dict[str, Any]

class RetrievalResponse(BaseModel):
    records: List[Record]

class ErrorResponse(BaseModel):
    error_code: int
    error_msg: str

# --- FastAPI App ---
app = FastAPI()

# --- API Key Authentication ---
API_KEY = os.getenv("DIFY_API_KEY", "your_dify_secret_key")

@app.middleware("http")
async def verify_api_key(request: Request, call_next):
    if request.url.path == "/retrieval":
        auth_header = request.headers.get("Authorization")
        if not auth_header:
            return JSONResponse(
                status_code=status.HTTP_401_UNAUTHORIZED,
                content={"error_code": 1001, "error_msg": "Authorization header is missing."}
            )
        parts = auth_header.split()
        if len(parts) != 2 or parts[0].lower() != "bearer":
            return JSONResponse(
                status_code=status.HTTP_401_UNAUTHORIZED,
                content={"error_code": 1001, "error_msg": "Invalid Authorization header format. Expected format is 'Bearer <api-key>'."}
            )
        token = parts[1]
        if token != API_KEY:
            return JSONResponse(
                status_code=status.HTTP_403_FORBIDDEN,
                content={"error_code": 1002, "error_msg": "Authorization failed."}
            )
    response = await call_next(request)
    return response

@app.post("/retrieval", response_model=RetrievalResponse, responses={404: {"model": ErrorResponse}})
async def retrieval(request: RetrievalRequest):
    knowledge_id = request.knowledge_id
    db_path = os.path.join(DATA_DIR, f"{knowledge_id}_meta.db")
    index_path = os.path.join(DATA_DIR, f"{knowledge_id}_index.faiss")

    if not os.path.exists(db_path) or not os.path.exists(index_path):
        return JSONResponse(status_code=404, content={"error_code": 2001, "error_msg": f"Knowledge base '{knowledge_id}' does not exist."})

    try:
        index = faiss.read_index(index_path)
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to load knowledge base artifacts: {e}")

    queries = [request.query]
    try:
        system_prompt = """You are a query expansion assistant. Based on the user query, generate 3 alternative, keyword-focused search queries for a technical programming language knowledge base. Return a JSON list of strings. For example, for the query 'how to use create apply', you could generate ["LPC apply create function", "documentation for apply create", "create apply"]."""
        completion = client.chat.completions.create(
            model="qwen-flash",
            messages=[{'role': 'system', 'content': system_prompt}, {'role': 'user', 'content': request.query}]
        )
        expanded_queries = json.loads(completion.choices[0].message.content)
        queries.extend(expanded_queries)
    except Exception as e:
        # In production, we can log this error to a file instead of printing
        pass

    candidate_ids = set()
    for q in queries:
        try:
            response = client.embeddings.create(model="text-embedding-v4", input=q)
            query_vector = np.array([response.data[0].embedding], dtype='float32')
            distances, indices = index.search(query_vector, 20)
            for idx in indices[0]:
                if idx != -1: candidate_ids.add(int(idx) + 1)
        except Exception as e:
            # In production, we can log this error to a file
            pass

    if not candidate_ids:
        conn.close()
        return RetrievalResponse(records=[])

    placeholders = ', '.join('?' for _ in candidate_ids)
    cursor.execute(f"SELECT * FROM chunks WHERE chunk_id IN ({placeholders})", tuple(candidate_ids))
    candidate_rows = cursor.fetchall()

    doc_texts_for_rerank = [row['chunk_text'] for row in candidate_rows]

    try:
        rerank_resp = dashscope.TextReRank.call(
            model="gte-rerank-v2",
            query=request.query,
            documents=doc_texts_for_rerank,
            return_documents=False
        )
    except Exception as e:
        conn.close()
        raise HTTPException(status_code=500, detail=f"Re-rank API call failed: {e}")

    records = []
    if rerank_resp.status_code == HTTPStatus.OK:
        score_threshold = request.retrieval_setting.score_threshold
        top_k_results = rerank_resp.output.results[:request.retrieval_setting.top_k]

        for result in top_k_results:
            if result.relevance_score >= score_threshold:
                original_doc_info = candidate_rows[result.index]
                metadata = json.loads(original_doc_info['metadata_json'])
                records.append(
                    Record(
                        content=original_doc_info['chunk_text'],
                        score=result.relevance_score,
                        title=metadata.get('topic', os.path.basename(original_doc_info['file_path'])),
                        metadata=metadata
                    )
                )
    else:
        # In production, we can log this error to a file
        pass

    conn.close()
    return RetrievalResponse(records=records)

if __name__ == "__main__":
    import uvicorn
    print("Starting API server...")
    print(f"To test, send a POST request to http://127.0.0.1:8080/retrieval")
    uvicorn.run(app, host="0.0.0.0", port=8080)