Chat: update response

pyproject.toml

@@ -23,7 +23,7 @@ sse-starlette = "^2.0.0"
 boto3 = "^1.0.0"
 httpx = "^0.27.0"
 chromadb = "^0.4.16"
-cohere = "^5.4.0"
+cohere = "^5.5.6"
 llama-index = "^0.10.11"
 inquirer = "^3.2.4"
 langchain-community = "^0.0.32"
@@ -75,7 +75,7 @@ optional = true
 
 [tool.poetry.group.community.dependencies]
 wolframalpha = "^5.0.0"
-transformers = "^4.40.1"
+transformers = "^4.3.3"
 torch = "^2.3.0"
 llama-cpp-python = "^0.2.67"
 llama-index = "^0.10.11"

src/backend/schemas/chat.py

@@ -17,6 +17,7 @@ class ChatRole(StrEnum):
     CHATBOT = "CHATBOT"
     USER = "USER"
     SYSTEM = "SYSTEM"
+    TOOL = "TOOL"
 
 
 class ChatCitationQuality(StrEnum):
@@ -39,9 +40,13 @@ class ChatMessage(BaseModel):
     role: ChatRole = Field(
         title="One of CHATBOT|USER|SYSTEM to identify who the message is coming from.",
     )
-    message: str = Field(
+    message: str | None = Field(
         title="Contents of the chat message.",
     )
+    tool_results: List[Dict[str, Any]] | None = Field(
+        title="Results from the tool call.",
+        default=[],
+    )
 
     def to_dict(self) -> Dict[str, str]:
         return {"role": self.role, "message": self.message}
@@ -50,9 +55,6 @@ def to_dict(self) -> Dict[str, str]:
 # TODO: fix titles of these types
 class ChatResponse(BaseModel):
     event_type: ClassVar[StreamEvent] = Field()
-    is_finished: bool = Field(
-        title="Denotes whether or not the chat stream has finished.",
-    )
 
 
 class StreamStart(ChatResponse):
@@ -150,7 +152,6 @@ class StreamToolCallsGeneration(ChatResponse):
 class StreamEnd(ChatResponse):
     response_id: str | None = Field(default=None)
     event_type: ClassVar[StreamEvent] = StreamEvent.STREAM_END
-    is_finished: ClassVar[bool] = True
     generation_id: str | None = Field(default=None)
     conversation_id: str | None = Field(default=None)
     text: str = Field(

src/backend/services/chat.py

@@ -549,8 +549,9 @@ def generate_chat_response(
 
         chat_history.append(
             ChatMessage(
-                role=message["role"],
-                message=message["message"],
+                role=message.get("role", ChatRole.USER),
+                message=message.get("message"),
+                tool_results=message.get("tool_results"),
             )
         )
 
@@ -587,7 +588,6 @@ def generate_chat_response(
         documents=documents,
         search_results=response.get("search_results", []),
         event_type=StreamEvent.NON_STREAMED_CHAT_RESPONSE,
-        is_finished=True,
         conversation_id=conversation_id,
         tool_calls=tool_calls,
     )
@@ -624,7 +624,6 @@ def generate_langchain_chat_stream(
             ChatResponseEvent(
                 event=StreamEvent.STREAM_START,
                 data=StreamStart(
-                    is_finished=False,
                     conversation_id=conversation_id,
                 ),
             )
@@ -661,7 +660,6 @@ def generate_langchain_chat_stream(
 
                     # shape: "Plan: I will search for tips on writing an essay and fun facts about the Roman Empire. I will then write an answer using the information I find.\nAction: ```json\n[\n    {\n        \"tool_name\": \"internet_search\",\n        \"parameters\": {\n            \"query\": \"tips for writing an essay\"\n        }\n    },\n    {\n        \"tool_name\": \"internet_search\",\n        \"parameters\": {\n            \"query\": \"fun facts about the roman empire\"\n        }\n
                     stream_event = StreamToolInput(
-                        is_finished=False,
                         # TODO: switch to diff types
                         input_type=ToolInputType.CODE,
                         tool_name=tool_name,
@@ -690,7 +688,6 @@ def generate_langchain_chat_stream(
                 if isinstance(result, list):
                     stream_event = StreamToolResult(
                         tool_name=step.action.tool,
-                        is_finished=False,
                         result=result,
                         documents=[],
                     )
@@ -708,7 +705,6 @@ def generate_langchain_chat_stream(
                 if isinstance(result, dict):
                     stream_event = StreamToolResult(
                         tool_name=step.action.tool,
-                        is_finished=False,
                         result=result,
                         documents=[],
                     )

src/backend/tests/model_deployments/mock_deployments/mock_azure.py

@@ -59,16 +59,13 @@ def invoke_chat_stream(
             {
                 "event_type": StreamEvent.STREAM_START,
                 "generation_id": "test",
-                "is_finished": False,
             },
             {
                 "event_type": StreamEvent.TEXT_GENERATION,
                 "text": "This is a test.",
-                "is_finished": True,
             },
             {
                 "event_type": StreamEvent.STREAM_END,
-                "is_finished": True,
                 "generation_id": "test",
                 "citations": [],
                 "documents": [],

src/backend/tests/model_deployments/mock_deployments/mock_bedrock.py

@@ -57,16 +57,13 @@ def invoke_chat_stream(
             {
                 "event_type": StreamEvent.STREAM_START,
                 "generation_id": "test",
-                "is_finished": False,
             },
             {
                 "event_type": StreamEvent.TEXT_GENERATION,
                 "text": "This is a test.",
-                "is_finished": True,
             },
             {
                 "event_type": StreamEvent.STREAM_END,
-                "is_finished": True,
                 "generation_id": "test",
                 "citations": [],
                 "documents": [],

src/backend/tests/model_deployments/mock_deployments/mock_cohere_platform.py

@@ -57,16 +57,13 @@ def invoke_chat_stream(
             {
                 "event_type": StreamEvent.STREAM_START,
                 "generation_id": "test",
-                "is_finished": False,
             },
             {
                 "event_type": StreamEvent.TEXT_GENERATION,
                 "text": "This is a test.",
-                "is_finished": True,
             },
             {
                 "event_type": StreamEvent.STREAM_END,
-                "is_finished": True,
                 "generation_id": "test",
                 "citations": [],
                 "documents": [],

src/backend/tests/model_deployments/mock_deployments/mock_sagemaker.py

@@ -31,16 +31,13 @@ def invoke_chat_stream(
             {
                 "event_type": StreamEvent.STREAM_START,
                 "generation_id": "test",
-                "is_finished": False,
             },
             {
                 "event_type": StreamEvent.TEXT_GENERATION,
                 "text": "This is a test.",
-                "is_finished": True,
             },
             {
                 "event_type": StreamEvent.STREAM_END,
-                "is_finished": True,
                 "generation_id": "test",
                 "citations": [],
                 "documents": [],

src/community/model_deployments/hugging_face.py

@@ -72,20 +72,17 @@ def invoke_chat_stream(self, chat_request: CohereChatRequest, **kwargs: Any) ->
         yield {
             "event-type": "event-start",
             "generation_id": "",
-            "is_finished": False,
         }
 
         gen_text = self.invoke_chat(chat_request)
 
         yield {
             "event-type": "text-generation",
             "text": gen_text.get("text", ""),
-            "is_finished": False,
         }
 
         yield {
             "event-type": "stream-end",
-            "is_finished": True,
             "finish_reason": "COMPLETE",
         }
 

src/community/model_deployments/local_model.py

@@ -49,20 +49,17 @@ def invoke_chat_stream(self, chat_request: CohereChatRequest, **kwargs: Any) ->
         yield {
             "event_type": "stream-start",
             "generation_id": "",
-            "is_finished": False,
         }
 
         for item in stream:
             yield {
                 "event_type": "text-generation",
                 "text": item["choices"][0]["text"],
-                "is_finished": False,
             }
 
         yield {
             "event_type": "stream-end",
             "finish_reason": "COMPLETE",
-            "is_finished": True,
         }
 
     def invoke_chat(self, chat_request: CohereChatRequest, **kwargs: Any) -> Any: