input sanitizing (langchain-ai#211)

main.py

@@ -2,11 +2,13 @@
 import asyncio
 import os
 from operator import itemgetter
-from typing import Dict, List, Optional, Sequence
+from typing import Dict, List, Optional, Sequence, Union
+from uuid import UUID
 
 import langsmith
 import weaviate
-from fastapi import FastAPI, Request
+from constants import WEAVIATE_DOCS_INDEX_NAME
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings
@@ -27,8 +29,6 @@
 from langserve import add_routes
 from langsmith import Client
 from pydantic import BaseModel
-
-from constants import WEAVIATE_DOCS_INDEX_NAME
 from voyage import VoyageEmbeddings
 
 RESPONSE_TEMPLATE = """\
@@ -218,34 +218,45 @@ def create_chain(
 add_routes(app, answer_chain, path="/chat", input_type=ChatRequest)
 
 
+class SendFeedbackBody(BaseModel):
+    run_id: UUID
+    key: str = "user_score"
+
+    score: Union[float, int, bool, None] = None
+    feedback_id: Optional[UUID] = None
+    comment: Optional[str] = None
+
+
 @app.post("/feedback")
-async def send_feedback(request: Request):
-    data = await request.json()
-    run_id = data.get("run_id")
-    if run_id is None:
-        return {
-            "result": "No LangSmith run ID provided",
-            "code": 400,
-        }
-    key = data.get("key", "user_score")
-    vals = {**data, "key": key}
-    client.create_feedback(**vals)
+async def send_feedback(body: SendFeedbackBody):
+    client.create_feedback(
+        body.run_id,
+        body.key,
+        score=body.score,
+        comment=body.comment,
+        feedback_id=body.feedback_id,
+    )
     return {"result": "posted feedback successfully", "code": 200}
 
 
+class UpdateFeedbackBody(BaseModel):
+    feedback_id: UUID
+    score: Union[float, int, bool, None] = None
+    comment: Optional[str] = None
+
+
 @app.patch("/feedback")
-async def update_feedback(request: Request):
-    data = await request.json()
-    feedback_id = data.get("feedback_id")
+async def update_feedback(body: UpdateFeedbackBody):
+    feedback_id = body.feedback_id
     if feedback_id is None:
         return {
             "result": "No feedback ID provided",
             "code": 400,
         }
     client.update_feedback(
         feedback_id,
-        score=data.get("score"),
-        comment=data.get("comment"),
+        score=body.score,
+        comment=body.comment,
     )
     return {"result": "patched feedback successfully", "code": 200}
 
@@ -268,16 +279,19 @@ async def aget_trace_url(run_id: str) -> str:
     return await _arun(client.share_run, run_id)
 
 
+class GetTraceBody(BaseModel):
+    run_id: UUID
+
+
 @app.post("/get_trace")
-async def get_trace(request: Request):
-    data = await request.json()
-    run_id = data.get("run_id")
+async def get_trace(body: GetTraceBody):
+    run_id = body.run_id
     if run_id is None:
         return {
             "result": "No LangSmith run ID provided",
             "code": 400,
         }
-    return await aget_trace_url(run_id)
+    return await aget_trace_url(str(run_id))
 
 
 if __name__ == "__main__":