typing fixes

Author: adnaans

vocode/streaming/agent/anthropic_utils.py

@@ -29,7 +29,7 @@ def format_anthropic_chat_messages_from_transcript(
     # TODO: reliably count tokens of Anthropic messages so that we don't exceed the context window
 
 def merge_bot_messages_for_langchain(messages: list[tuple]) -> list[tuple]:
-    merged_messages = []
+    merged_messages: list[tuple] = []
     for role, message in messages:
         if role == "ai" and merged_messages and merged_messages[-1][0] == "ai":
             merged_messages[-1] = ("ai", merged_messages[-1][1] + message)

vocode/streaming/agent/langchain_agent.py

@@ -3,7 +3,7 @@
 import sentry_sdk
 from loguru import logger
 
-from langchain_core.messages.ai import AIMessageChunk
+from langchain_core.messages.base import BaseMessage as LangchainBaseMessage
 from langchain_core.prompts import ChatPromptTemplate
 from langchain.chat_models import init_chat_model
 
@@ -44,10 +44,13 @@ def __init__(
 
     async def token_generator(
         self,
-        gen: AsyncIterator[AIMessageChunk],
-    ) -> AsyncGenerator[str | FunctionFragment, None]:
+        gen: AsyncIterator[LangchainBaseMessage],
+    ) -> AsyncGenerator[str, None]:
         async for chunk in gen:
-            yield chunk.content
+            if isinstance(chunk.content, str):
+                yield chunk.content
+            else:
+                raise ValueError(f"Received unexpected message type {type(chunk)} from Langchain. Expected str.")
 
     def format_langchain_messages_from_transcript(self) -> list[tuple]:
         if not self.transcript: