cognee/level_3/api.py

import logging
import os
from typing import Dict, Any

import uvicorn
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from level_2_pdf_vectorstore__dlt_contracts import Memory
from dotenv import load_dotenv
# Set up logging
logging.basicConfig(
    level=logging.INFO,  # Set the logging level (e.g., DEBUG, INFO, WARNING, ERROR, CRITICAL)
    format="%(asctime)s [%(levelname)s] %(message)s",  # Set the log message format
)

logger = logging.getLogger(__name__)



load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
app = FastAPI(debug=True)

from auth.cognito.JWTBearer import JWTBearer
from auth.auth import jwks
auth = JWTBearer(jwks)

from fastapi import Depends
class ImageResponse(BaseModel):
    success: bool
    message: str

@app.get(
    "/",
)
async def root():
    """
    Root endpoint that returns a welcome message.
    """
    return {"message": "Hello, World, I am alive!"}

@app.get("/health")
def health_check():
    """
    Health check endpoint that returns the server status.
    """
    return {"status": "OK"}
class Payload(BaseModel):
    payload: Dict[str, Any]

def memory_factory(memory_type):
    load_dotenv()

    class Payload(BaseModel):
        payload: Dict[str, Any]

    @app.post("/{memory_type}/add-memory", response_model=dict)
    async def add_memory(
        payload: Payload,
        # files: List[UploadFile] = File(...),
    ):
        try:
            logging.info(" Init PDF processing")

            decoded_payload = payload.payload

            Memory_ = Memory(user_id=decoded_payload["user_id"])

            await Memory_.async_init()

            memory_class = getattr(Memory_, f"_add_{memory_type}_memory", None)
            output = await memory_class(
                observation=decoded_payload["prompt"],
                loader_settings=decoded_payload["loader_settings"],
                params=decoded_payload["params"],
            )
            return JSONResponse(content={"response": output}, status_code=200)

        except Exception as e:
            return JSONResponse(
                content={"response": {"error": str(e)}}, status_code=503
            )

    @app.post("/{memory_type}/fetch-memory", response_model=dict)
    async def fetch_memory(
        payload: Payload,
        # files: List[UploadFile] = File(...),
    ):
        try:
            decoded_payload = payload.payload

            Memory_ = Memory(user_id=decoded_payload["user_id"])

            await Memory_.async_init()

            memory_class = getattr(Memory_, f"_fetch_{memory_type}_memory", None)
            output = memory_class(observation=decoded_payload["prompt"])
            return JSONResponse(content={"response": output}, status_code=200)

        except Exception as e:
            return JSONResponse(
                content={"response": {"error": str(e)}}, status_code=503
            )

    @app.post("/{memory_type}/delete-memory", response_model=dict)
    async def delete_memory(
        payload: Payload,
        # files: List[UploadFile] = File(...),
    ):
        try:
            decoded_payload = payload.payload

            Memory_ = Memory(user_id=decoded_payload["user_id"])

            await Memory_.async_init()

            memory_class = getattr(Memory_, f"_delete_{memory_type}_memory", None)
            output = memory_class(observation=decoded_payload["prompt"])
            return JSONResponse(content={"response": output}, status_code=200)

        except Exception as e:
            return JSONResponse(
                content={"response": {"error": str(e)}}, status_code=503
            )


memory_list = ["episodic", "buffer", "semantic"]
for memory_type in memory_list:
    memory_factory(memory_type)


@app.get("/available-buffer-actions", response_model=dict)
async def available_buffer_actions(
    payload: Payload,
    # files: List[UploadFile] = File(...),
):
    try:
        decoded_payload = payload.payload

        Memory_ = Memory(user_id=decoded_payload["user_id"])

        await Memory_.async_init()

        # memory_class = getattr(Memory_, f"_delete_{memory_type}_memory", None)
        output = await Memory_._available_operations()
        return JSONResponse(content={"response": output}, status_code=200)

    except Exception as e:
        return JSONResponse(content={"response": {"error": str(e)}}, status_code=503)


@app.post("/run-buffer", response_model=dict)
async def run_buffer(
    payload: Payload,
    # files: List[UploadFile] = File(...),
):
    try:
        decoded_payload = payload.payload

        Memory_ = Memory(user_id=decoded_payload["user_id"])

        await Memory_.async_init()

        # memory_class = getattr(Memory_, f"_delete_{memory_type}_memory", None)
        output = await Memory_._run_main_buffer(
            user_input=decoded_payload["prompt"], params=decoded_payload["params"], attention_modulators=decoded_payload["attention_modulators"]
        )
        return JSONResponse(content={"response": output}, status_code=200)

    except Exception as e:
        return JSONResponse(content={"response": {"error": str(e)}}, status_code=503)


@app.post("/buffer/create-context", response_model=dict)
async def create_context(
    payload: Payload,
    # files: List[UploadFile] = File(...),
):
    try:
        decoded_payload = payload.payload

        Memory_ = Memory(user_id=decoded_payload["user_id"])

        await Memory_.async_init()

        # memory_class = getattr(Memory_, f"_delete_{memory_type}_memory", None)
        output = await Memory_._create_buffer_context(
            user_input=decoded_payload["prompt"], params=decoded_payload["params"], attention_modulators=decoded_payload["attention_modulators"]
        )
        return JSONResponse(content={"response": output}, status_code=200)

    except Exception as e:
        return JSONResponse(content={"response": {"error": str(e)}}, status_code=503)


@app.post("/buffer/get-tasks", response_model=dict)
async def create_context(
    payload: Payload,
    # files: List[UploadFile] = File(...),
):
    try:
        decoded_payload = payload.payload

        Memory_ = Memory(user_id=decoded_payload["user_id"])

        await Memory_.async_init()

        # memory_class = getattr(Memory_, f"_delete_{memory_type}_memory", None)
        output = await Memory_._get_task_list(
            user_input=decoded_payload["prompt"], params=decoded_payload["params"], attention_modulators=decoded_payload["attention_modulators"]
        )
        return JSONResponse(content={"response": output}, status_code=200)

    except Exception as e:
        return JSONResponse(content={"response": {"error": str(e)}}, status_code=503)


@app.post("/buffer/provide-feedback", response_model=dict)
async def provide_feedback(
    payload: Payload,
    # files: List[UploadFile] = File(...),
):
    try:
        decoded_payload = payload.payload

        Memory_ = Memory(user_id=decoded_payload["user_id"])

        await Memory_.async_init()

        # memory_class = getattr(Memory_, f"_delete_{memory_type}_memory", None)
        if decoded_payload["total_score"] is None:

            output = await Memory_._provide_feedback(
                user_input=decoded_payload["prompt"], params=decoded_payload["params"], attention_modulators=None, total_score=decoded_payload["total_score"]
            )
            return JSONResponse(content={"response": output}, status_code=200)
        else:
            output = await Memory_._provide_feedback(
                user_input=decoded_payload["prompt"], params=decoded_payload["params"], attention_modulators=decoded_payload["attention_modulators"], total_score=None
            )
            return JSONResponse(content={"response": output}, status_code=200)


    except Exception as e:
        return JSONResponse(content={"response": {"error": str(e)}}, status_code=503)
def start_api_server(host: str = "0.0.0.0", port: int = 8000):
    """
    Start the API server using uvicorn.

    Parameters:
    host (str): The host for the server.
    port (int): The port for the server.
    """
    try:
        logger.info(f"Starting server at {host}:{port}")
        uvicorn.run(app, host=host, port=port)
    except Exception as e:
        logger.exception(f"Failed to start server: {e}")
        # Here you could add any cleanup code or error recovery code.


if __name__ == "__main__":
    start_api_server()