feat: solution for 'Агент с RAG-памятью'

src/main.py

@@ -1,21 +1,44 @@
-#!/usr/bin/env python3
 """
-Main entry point for the knowledge‑base agent.
+FastAPI application exposing the RAG agent as a REST endpoint.
+
+This file is optional but useful for running the agent in a container.
 """
 
-from .knowledge_base import KnowledgeBase
-from .tools.knowledge_base_tool import KnowledgeBaseTool
-from .cli import run_cli
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+from .agent import RAGAgent
+
+app = FastAPI(title="RAG Agent API")
+
+# Initialize a global agent instance
+agent = RAGAgent(model="llama2", top_k=3)
 
 
-def main() -> None:
+class Document(BaseModel):
+    text: str
+
+
+class Query(BaseModel):
+    query: str
+
+
+@app.post("/documents")
+def add_document(doc: Document):
     """
-    Create the knowledge base, wrap it in a tool, and start the CLI.
+    Add a document to the agent's memory.
     """
-    kb = KnowledgeBase()
-    kb_tool = KnowledgeBaseTool(kb)
-    run_cli(kb_tool)
+    agent.add_document(doc.text)
+    return {"status": "added"}
 
 
-if __name__ == "__main__":
-    main()
+@app.post("/ask")
+def ask(query: Query):
+    """
+    Get an answer to a query using the RAG agent.
+    """
+    try:
+        answer = agent.get_response(query.query)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc))
+    return {"answer": answer}