ic

apps/riesgos/api/server/v1.py

@@ -0,0 +1,64 @@
+import uuid
+from typing import Literal
+
+from fastapi import APIRouter
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+from api import config, services
+from api.agent import MayaRiesgos
+
+router = APIRouter(prefix="/api/v1")
+agent = MayaRiesgos()
+
+
+class Message(BaseModel):
+    conversation_id: uuid.UUID
+    prompt: str
+
+
+@router.post("/conversation")
+async def create_conversation():
+    conversation_id = uuid.uuid4()
+    await services.create_conversation(conversation_id, agent.system_prompt)
+    return {"conversation_id": conversation_id}
+
+
+@router.post("/message")
+async def send(message: Message, stream: bool = False):
+    if stream is True:
+
+        def b64_sse(func):
+            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")
+    else:
+        response = await services.generate(
+            agent, message.prompt, message.conversation_id
+        )
+        return response
+
+
+class Feedback(BaseModel):
+    key: str
+    rating: Literal["Good", "Bad", "None"]
+
+
+@router.post("/feedback")
+async def register_feedback(data: Feedback):
+    if data.rating:
+        langfuse = config.get_langfuse()
+        langfuse.score(
+            id=data.key + "-rating",
+            trace_id=data.key,
+            name="Rating",
+            value=data.rating,
+        )