Enhance the API documentation. (#54)

* feat: add FeedbackResponse schema for api docs

* feat: add sse response for api docs

* chore: add sse DONE example

src/schema/__init__.py

@@ -2,6 +2,7 @@
     AgentResponse,
     ChatMessage,
     Feedback,
+    FeedbackResponse,
     StreamInput,
     UserInput,
     convert_message_content_to_string,
@@ -13,5 +14,6 @@
     "ChatMessage",
     "StreamInput",
     "Feedback",
+    "FeedbackResponse",
     "convert_message_content_to_string",
 ]

src/schema/schema.py

@@ -169,3 +169,7 @@ class Feedback(BaseModel):
         default={},
         examples=[{"comment": "In-line human feedback"}],
     )
+
+
+class FeedbackResponse(BaseModel):
+    status: Literal["success"] = "success"

src/service/service.py

@@ -6,7 +6,7 @@
 from typing import Any
 from uuid import uuid4
 
-from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi import FastAPI, HTTPException, Request, Response, status
 from fastapi.responses import StreamingResponse
 from langchain_core._api import LangChainBetaWarning
 from langchain_core.runnables import RunnableConfig
@@ -15,7 +15,14 @@
 from langsmith import Client as LangsmithClient
 
 from agent import research_assistant
-from schema import ChatMessage, Feedback, StreamInput, UserInput, convert_message_content_to_string
+from schema import (
+    ChatMessage,
+    Feedback,
+    FeedbackResponse,
+    StreamInput,
+    UserInput,
+    convert_message_content_to_string,
+)
 
 warnings.filterwarnings("ignore", category=LangChainBetaWarning)
 
@@ -142,7 +149,21 @@ async def message_generator(user_input: StreamInput) -> AsyncGenerator[str, None
     yield "data: [DONE]\n\n"
 
 
-@app.post("/stream")
+def _sse_response_example() -> dict[str, Any]:
+    return {
+        status.HTTP_200_OK: {
+            "description": "Server Sent Event Response",
+            "content": {
+                "text/event-stream": {
+                    "example": "data: {'type': 'token', 'content': 'Hello'}\n\ndata: {'type': 'token', 'content': ' World'}\n\ndata: [DONE]\n\n",
+                    "schema": {"type": "string"},
+                }
+            },
+        }
+    }
+
+
+@app.post("/stream", response_class=StreamingResponse, responses=_sse_response_example())
 async def stream_agent(user_input: StreamInput) -> StreamingResponse:
     """
     Stream the agent's response to a user input, including intermediate messages and tokens.
@@ -154,7 +175,7 @@ async def stream_agent(user_input: StreamInput) -> StreamingResponse:
 
 
 @app.post("/feedback")
-async def feedback(feedback: Feedback) -> dict[str, str]:
+async def feedback(feedback: Feedback) -> FeedbackResponse:
     """
     Record feedback for a run to LangSmith.
 
@@ -170,4 +191,4 @@ async def feedback(feedback: Feedback) -> dict[str, str]:
         score=feedback.score,
         **kwargs,
     )
-    return {"status": "success"}
+    return FeedbackResponse()