ic

.templates/rag/{{project_name}}/api/server.py

@@ -0,0 +1,49 @@
+import uuid
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+from api import services
+from api.agent import Agent
+from api.config import config
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI):
+    await config.init_mongo_db()
+    yield
+
+
+app = FastAPI(lifespan=lifespan)
+agent = Agent()
+
+
+@app.post("/api/v1/conversation")
+async def create_conversation():
+    conversation_id = uuid.uuid4()
+    await services.create_conversation(conversation_id, agent.system_prompt)
+    return {"conversation_id": conversation_id}
+
+
+class Message(BaseModel):
+    conversation_id: uuid.UUID
+    prompt: str
+
+@app.post("/api/v1/message")
+async def send(message: Message):
+    def b64_sse(func):
+        """Este helper transforma un generador de strings a un generador del protocolo SSE"""
+
+        async def wrapper(*args, **kwargs):
+            async for chunk in func(*args, **kwargs):
+                content = chunk.model_dump_json()
+                data = f"data: {content}\n\n"
+                yield data
+
+        return wrapper
+
+    sse_stream = b64_sse(services.stream)
+    generator = sse_stream(agent, message.prompt, message.conversation_id)
+    return StreamingResponse(generator, media_type="text/event-stream")