fast cgi daemon

import os, json
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

app = FastAPI()

import logging
logging.basicConfig(level=logging.DEBUG)

# CORS for OpenWebUI on 192.168.88.239
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://192.168.88.239"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class QueryRequest(BaseModel):
    question: str
    top_k: int = 3
    system_prompt: str = ""

chunks = []
chunk_ids = []
vectorizer = None
embeddings = None


def load_docs_for_tenant(tenant_id):
    global chunks, chunk_ids, vectorizer, embeddings
    chunks.clear()
    chunk_ids.clear()

    # Look for JSON files in the docling.json directory
    docling_dir = f"./docling/{tenant_id}/docling.json"
    if not os.path.exists(docling_dir):
        print(f"Warning: No docling directory found for tenant {tenant_id} at {docling_dir}")
        # Initialize empty vectorizer to prevent errors
        vectorizer = TfidfVectorizer()
        embeddings = None
        return

    # If it's a directory, look for JSON files inside it
    if os.path.isdir(docling_dir):
        json_files = [f for f in os.listdir(docling_dir) if f.endswith('.json')]
        if not json_files:
            print(f"Warning: No JSON files found in {docling_dir}")
            vectorizer = TfidfVectorizer()
            embeddings = None
            return

        # Load all JSON files in the directory
        for json_file in json_files:
            path = os.path.join(docling_dir, json_file)
            try:
                with open(path, "r", encoding="utf-8") as f:
                    data = json.load(f)

                    # Extract text from the texts section
                    if isinstance(data, dict) and "texts" in data:
                        texts = data["texts"]
                        for i, text_item in enumerate(texts):
                            if isinstance(text_item, dict) and "text" in text_item:
                                text_content = text_item["text"].strip()
                                if text_content:
                                    chunks.append(text_content)
                                    chunk_ids.append(f"{json_file}_text_{i}")

                    # Also try to extract from body children if they have direct text
                    elif isinstance(data, dict) and "body" in data:
                        body_data = data["body"]
                        if isinstance(body_data, dict) and "children" in body_data:
                            # Process children recursively to extract text
                            def extract_text_from_children(children):
                                text_parts = []
                                for child in children:
                                    if isinstance(child, dict):
                                        if "text" in child:
                                            text_parts.append(child["text"])
                                        if "children" in child:
                                            text_parts.extend(extract_text_from_children(child["children"]))
                                return text_parts

                            text_parts = extract_text_from_children(body_data["children"])
                            for i, text in enumerate(text_parts):
                                if text.strip():
                                    chunks.append(text.strip())
                                    chunk_ids.append(f"{json_file}_chunk_{i}")
            except Exception as e:
                print(f"Error loading {path}: {e}")
                continue
    else:
        # Fallback to original behavior if it's a file
        try:
            with open(docling_dir, "r", encoding="utf-8") as f:
                data = json.load(f)
                for i, entry in enumerate(data):
                    body = entry.get("body", "").strip()
                    if body:
                        chunks.append(body)
                        chunk_ids.append(entry.get("id", f"chunk{i}"))
        except Exception as e:
            print(f"Error loading {docling_dir}: {e}")

    if chunks:
        print(f"Loaded {len(chunks)} chunks for tenant {tenant_id}")
        vectorizer = TfidfVectorizer().fit(chunks)
        embeddings = vectorizer.transform(chunks)
    else:
        print(f"Warning: No valid chunks found for tenant {tenant_id}")
        vectorizer = TfidfVectorizer()
        embeddings = None


def query(question: str, top_k: int = 3):
    if not chunks or embeddings is None:
        return []

    q_vec = vectorizer.transform([question])
    sims = cosine_similarity(q_vec, embeddings)[0]
    indices = sims.argsort()[-top_k:][::-1]
    return [{"id": chunk_ids[i], "text": chunks[i], "score": float(sims[i])} for i in indices]


@app.post("/query")
async def query_docs(req: QueryRequest, request: Request):
    # Extract tenant_id from system prompt
    tenant_id = "default"
    if req.system_prompt:
        import re
        # Look for "tenant: {tenant_id}" at the beginning of the prompt
        tenant_match = re.search(r'^tenant:\s*(\w+)', req.system_prompt.strip(), re.IGNORECASE | re.MULTILINE)
        if tenant_match:
            tenant_id = tenant_match.group(1).lower()

    load_docs_for_tenant(tenant_id)
    return {"results": query(req.question, req.top_k)}


@app.get("/health")
async def health_check():
    return {"status": "ok"}


if __name__ == "__main__":
    print("Starting FastAPI server...")
    import uvicorn
    print("Uvicorn imported successfully")
    uvicorn.run(app, host="0.0.0.0", port=5000)
    print("Server started")