main.py

import base64
import contextlib
import json
import logging
import os
from logging.handlers import RotatingFileHandler
from typing import Dict

import openai
from dotenv import load_dotenv
from fastapi import Body, FastAPI, File, HTTPException, Query, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam

from db import ChatMessage, ChatRequest, LogEntry, Message, QueryRequest, _profiles, get_session, vdb
from utils import add_flows_to_rag, jgrep, setup_logger

load_dotenv()
logger = setup_logger("main-server")
ENV = "prod"
app = FastAPI(
    title="Atdevs AI API Specification",
    servers=[{"url": "http://94.136.186.136:8000"}, {"url": "http://localhost:8000"}],
    docs_url="/docs1",
    redoc_url="/docs",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=5)

# app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/admin", StaticFiles(directory="admin/build", html=True), name="admin")

WSL_WIN_HOST_IP = "172.24.192.1"


def chat(msgs: list[Message], profile="external_fastest"):
    """
    Interact with various LLM models through a unified chat interface. Profile any of Literal["external_fastest", "external_smartest", "private_mini", "private_large"]
    """
    default_key = "SGAIS_AI_TEST"
    profiles = {
        "external_fastest": {"base_url": "https://api.cerebras.ai/v1/", "model": "llama-3.3-70b", "api_key": os.environ["CEREBRAS_API_KEY"]},
        "external_smartest": {"base_url": "https://models.inference.ai.azure.com/", "model": ["gpt-4o", "qwen2.5-coder"][0], "api_key": os.getenv("GITPAT")},
        "private_fast": {"base_url": f"http://{WSL_WIN_HOST_IP}:11434/v1/", "model": "deepseek-r1:1.5b", "api_key": default_key},
        "private_balanced": {"base_url": f"http://{WSL_WIN_HOST_IP}:11434/v1/", "model": "qwen2.5-coder", "api_key": default_key},
        "private_smart": {"base_url": f"http://{WSL_WIN_HOST_IP}:11434/v1/", "model": "qwen2.5-coder:14b-instruct-q6_K", "api_key": default_key},
        "remote": {"base_url": "http://94.136.186.136:11434/v1/", "model": ["qwen2.5:14b", "qwen2.5-coder:14b-instruct-q6_K"][1], "api_key": ""},
        # "private_balanced": {"base_url": f"http://{WSL_WIN_HOST_IP}:11434/v1/", "model": "marco-o1:7b-q8_0", "api_key": None},
        # "groq": {"base_url": "https://api.groq.com/openai/v1/", "model": "llama-3.1-70b-versatile", "api_key": os.environ["GROQ_API_KEY"]},
    }
    if profile not in profiles:
        raise ValueError("Unmatched AI Model Profile")
    p: dict = profiles[profile]
    model = p.pop("model")
    client = openai.OpenAI(**p)
    response = client.chat.completions.create(
        model=model,
        messages=[
            ChatCompletionUserMessageParam(role="user", content=m.content)
            if m.role == "user"
            else ChatCompletionSystemMessageParam(role="system", content=m.content)
            for m in msgs
        ],
        # response_format={"type": "json_object"},
        extra_body={"options": {"num_ctx": 1024 * 8, "use_mmap": False, "temperature": 0.1, "low_vram": True}} if ("11434" in p["base_url"]) else {},
    )

    return response.choices[0].message.content if response.choices else None


@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
    """
    Simple chat endpoint that handles conversation with AI models.

    Args:
        request: ChatRequest containing:
            - messages: List of chat messages with role and content
            - profile: Optional AI model profile to use

    Returns:
        JSON response containing the AI's reply
    """
    try:
        # Convert ChatMessage objects to Message objects
        messages = [Message(role=msg.role, content=msg.content) for msg in request.messages]

        # Get response from AI
        response = chat(messages, profile=request.profile if request.profile else "external_fastest")

        # Log the interaction
        with get_session() as session:
            log_entry = LogEntry(request=json.dumps([m.dict() for m in request.messages]), response=response, meta={"profile": request.profile}, type="chat")
            session.add(log_entry)
            session.commit()

        return {"response": response, "status": "success"}

    except Exception as e:
        logger.error(f"Error in chat endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing chat request: {str(e)}")


@app.post("/openapi-spec-generator")
async def gen_openapi_spec(query: QueryRequest):
    """
    Generate OpenAPI v3 specification from natural language description.

    Args:
        request: Query containing the API description

    Returns:
        Generated OpenAPI specification as JSON
    """
    print(query)
    system = Message(
        role="system",
        content="You are an OpenAPI specification generator. Generate OpenAPI v3 specification based on the provided description, including descriptions for all endpoints and parameters. The output should be in only valid JSON format and follow OpenAPI v3 standards. Output Example ```json\n{...}\n```",
    )
    messages = [system, Message(role="user", content=query.prompt)]
    try:
        response = chat(messages, profile=query.profile)

        if not response:
            raise Exception("No response received from AI")

        resp = jgrep(response)
        logger.info(f"Generated OpenAPI spec: {response}")
        with get_session() as session:
            log_entry = LogEntry(
                request=query.prompt,
                response=json.dumps(resp),
                meta={"profile": query.profile, "parsed_response": resp, "messages": [m.dict() for m in messages]},
                type="openapi_spec",
            )
            session.add(log_entry)
            session.commit()
            session.refresh(log_entry)  # Refresh to get the ID and other DB-generated values
        return resp
    except Exception as e:
        raise e
        logger.error(f"Error generating OpenAPI spec: {e}")
        return {"error": str(e)}


@app.post("/flow-generator")
async def gen_ipaas_ui_flow_seq(query: QueryRequest, limit: int = 5):
    """vv
    Generate integration flow based on natural language description.

    Args:
        request: Query describing the desired integration flow\n
        limit: Maximum number of flows to generate

    Returns:
        Generated flow specification as JSON
    """
    logger.info(f"Generating flow for query: {query}")
    r = vdb.query(query.prompt)
    if not r or not r.get("documents") or not r["documents"]:
        logger.warning("No documents found for query")
        r_text = ""
    else:
        r_text = "\n\n\n\n".join(r["documents"][0])

    doc = open("./docs/flows-doc-simplified.md").read()
    messages = [
        Message(
            role="system",
            content=f"Closely observe my reference docs, from which you will generate a similar output. before that you will give a small paragraph thinking 100 words explaining what your strategy will be. <reference>\n{r_text}</reference>",
        ),
        Message(role="user", content=f"DO NOT CHANGE STRUCTURE OF MY JSON, just give as no other output JUST JSON only. ```json\ncontent```"),
        Message(role="user", content=query.prompt),
    ]
    print(*map(lambda x: x.content, messages))
    try:
        response = chat(messages, profile=query.profile)
        try:
            out = response
        except Exception:
            out = response.get("message", {}).get("tool_calls", []) if isinstance(response, dict) else []

        out = jgrep(out)
        logger.info(f"Generated flow: {out}")
        # log to db
        with get_session() as session:
            log_entry = LogEntry(request=json.dumps(query.dict()), response=json.dumps(out), meta={"profile": query.profile}, type="flow")
            session.add(log_entry)
            session.commit()
        return out
    except Exception as e:
        # raise internal server error with appropriate ,essage
        return HTTPException(status_code=500, detail=f"Error generating flow: {e}")


@app.post("/postman-generator")
async def gen_tests_for_postman(
    openapi_spec=Body(...),
    profile: _profiles = Query("private_fast", description="Profile to use for generation, To passed as query param, as body is openapi spec"),
):
    """
    Generate Postman tests from OpenAPI specification.

    Args:
        openapi_spec: OpenAPI specification as string
        profile: AI model profile to use for generation

    Returns:
        Generated Postman test collection as JSON

    Raises:
        HTTPException: If the input OpenAPI specification is not a string
    """
    if not isinstance(openapi_spec, str):
        raise HTTPException(status_code=400, detail="Invalid input type for OpenAPI specification")

    system = Message(
        role="system",
        content="You are a test generator. Create comprehensive postman collection with comprehensive tests for my given openapi spec. The output should be in JSON format. write the content in script.",
    )
    messages = [system, Message(role="user", content=openapi_spec)]
    response = jgrep(chat(messages, profile=profile))  # Changed from local-deepseek to local
    return response


@app.post("/kb")
async def knowledgebase_add_one(request: Request, source: str = "kb", id: str = None):
    """Add a document to the knowledge base.

    Args:
        request: The request containing the document text
        source: Source of the document (defaults to "kb")
        id: Optional ID for the document (if not provided, will be generated from content)
    """
    import re, hashlib

    try:
        data = (await request.body()).decode("utf-8")
        # Normalize document by replacing spaces with tabs and removing special characters
        norm_doc = re.sub(r"[\t\r]", "", data.replace("    ", "\t"))

        # Generate or use provided document ID
        doc_id = id if id else hashlib.sha256(norm_doc.encode()).hexdigest()

        # Create metadata dict ensuring all values are valid types
        metadata = {"source": str(source), "logid": str(id) if id else None}
        # Remove None values from metadata
        metadata = {k: v for k, v in metadata.items() if v is not None}

        # Add document to vector database
        vdb.add(doc_id, norm_doc, metadata)

        # Return the added document
        return vdb.collection.get(ids=[doc_id])

    except Exception as e:
        logger.error(f"Error adding document: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error adding documents: {str(e)}")


@app.get("/kb/query")
async def knowledgebase_query_relevant(q: str, limit: int = 5, include_metadata: bool = True, strict: bool = False):
    print(f"q={q} | strict={strict}")
    try:
        results = vdb.collection.query(
            query_embeddings=[vdb.embed(q)],
            n_results=limit,
            include=["documents", "distances", "metadatas"] if include_metadata else ["documents", "distances"],
            where_document={"$contains": q} if strict else None,
        )

        # Format results to match /kb endpoint structure
        matches = []
        if results:
            for i in range(len(results["documents"][0])):
                # Generate a unique ID for each document ( assumming no existing IDs )

                # Create a document object with id, content, and optional metadata
                document = {
                    "id": results["ids"][0][i],
                    "doc": results["documents"][0][i],
                    "metadata": results["metadatas"][0][i] if include_metadata else {},
                    "distance": results["distances"][0][i] if results.get("distances") else 0,
                }

                matches.append(document)

        return matches

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error performing search: {e!s}")


@app.get("/kb")
async def knowledgebase_all():
    try:
        result = vdb.collection.get(limit=1000)
        resp = []
        if not result["documents"]:
            return resp

        for i, d, m in zip(result["ids"], result["documents"], result["metadatas"] if result.get("metadatas") else []):
            resp.append({"id": i, "doc": d, "metadata": m})

        return resp

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error retrieving document: {str(e)}")


@app.get("/kb/{doc_id}")
async def knowledgebase_get_one(doc_id: str):
    try:
        result = vdb.collection.get(ids=[doc_id])

        if not result["documents"]:
            raise HTTPException(status_code=404, detail="Document not found")

        return {"id": doc_id, "content": result["documents"][0], "metadata": result["metadatas"][0] if result["metadatas"] else None}

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error retrieving document: {str(e)}")


@app.delete("/kb/{doc_id}")
async def knowledgebase_del_one(doc_id: str):
    try:
        vdb.collection.delete(ids=[doc_id])
        return {"message": f"Document {doc_id} deleted successfully"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error deleting document: {str(e)}")


@app.get("/logs")
async def get_all_logs(limit=1000):
    """Retrieve system logs with optional limit."""
    _ = get_session().query(LogEntry).limit(limit).all()
    _.reverse()
    return _


# route to delete a log id
@app.delete("/logs/{log_id}")
async def delete_log(log_id: int):
    """Delete a specific log entry by ID."""
    with get_session() as session:
        log_entry = session.query(LogEntry).filter(LogEntry.id == log_id).first()
        if not log_entry:
            raise
        return "ok"


# backup route which zips folder databases/*
@app.get("/backup")
async def backup():
    """
    Backup the database and return a zip file.
    """
    import io, datetime
    import zipfile
    from pathlib import Path

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    zip_buffer = io.BytesIO()
    db_path = Path("databases")

    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
        for file in db_path.rglob("*"):
            if file.is_file():
                zf.write(file, file.relative_to(db_path))

    zip_buffer.seek(0)
    return StreamingResponse(zip_buffer, media_type="application/zip", headers={"Content-Disposition": f"attachment; filename=backup_{timestamp}.zip"})


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True, reload_delay=2, reload_excludes=[".venv/*"])