app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from llm import LargeLanguageModel, Message
from memory import Memory

app = FastAPI(title="Helena AI Server", version="1.0.0")

llm = LargeLanguageModel()
memory = Memory()

class ChatRequest(BaseModel):
    message: str
    username: str = "user"
    image: Optional[str] = None


class ChatResponse(BaseModel):
    response: str
    memories_used: int


class MemoryStats(BaseModel):
    total_memories: int
    recent_memories: list


@app.get("/")
async def root():
    """Health check endpoint."""
    return {
        "status": "ok",
        "message": "running",
        "endpoints": {
            "chat": "/chat",
            "stats": "/stats",
        }
    }


@app.post("/chat", response_model=ChatResponse)
async def chat_with_memory(request: ChatRequest):
    try:
        memories = memory.retrieve(request.message, top_k=5)
        context = memory.build_context(request.message, top_k=5)
        
        if context.strip():
            full_input = f"Context:\n{context}\n\nUser Message: {request.message}"
        else:
            full_input = request.message
        
        user_message = Message(
            username=request.username,
            content=full_input,
            system=False,
            image=request.image
        )
        
        response = llm.generate(user_message)
        
        memory.add(request.username, request.message)
        
        memory.add("assistant", response.content)
        
        return ChatResponse(
            response=response.content,
            memories_used=len(memories)
        )
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error processing chat: {str(e)}"
        )


@app.get("/stats", response_model=MemoryStats)
async def get_memory_stats():
    try:
        count = memory.count()
        recent = memory.recent(limit=5)
        
        recent_formatted = [
            {
                "role": mem["role"],
                "content": mem["content"][:100] + "..." if len(mem["content"]) > 100 else mem["content"]
            }
            for mem in recent
        ]
        
        return MemoryStats(
            total_memories=count,
            recent_memories=recent_formatted
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error getting stats: {str(e)}"
        )


@app.post("/clear")
async def clear_conversation():
    try:
        llm.stm = []
        memory.clear()
        return {"status": "ok", "message": "Conversation cleared"}
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error clearing conversation: {str(e)}"
        )


@app.on_event("shutdown")
async def shutdown_event():
    memory.close()


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)