diff --git a/src/schema/__init__.py b/src/schema/__init__.py index 2214434..de6106e 100644 --- a/src/schema/__init__.py +++ b/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", ] diff --git a/src/schema/schema.py b/src/schema/schema.py index 97de418..7eeb381 100644 --- a/src/schema/schema.py +++ b/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" diff --git a/src/service/service.py b/src/service/service.py index beb6a31..03754c9 100644 --- a/src/service/service.py +++ b/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()