make the Q&A demo work

Author: mme

examples/coagents-qa/agent/my_agent/agent.py

@@ -1,20 +1,50 @@
 """Test Q&A Agent"""
 
+import os
 from typing import Any, cast
 from langchain_openai import ChatOpenAI
+from langchain_anthropic import ChatAnthropic
 from langgraph.graph import StateGraph, END
 from langgraph.graph import MessagesState
 from langgraph.checkpoint.memory import MemorySaver
 from langchain_core.runnables import RunnableConfig
-from langchain_core.messages import SystemMessage, ToolMessage
+from langchain_core.messages import HumanMessage, ToolMessage
 from copilotkit.langchain import (
   copilotkit_customize_config, copilotkit_exit, copilotkit_emit_message
 )
+from pydantic import BaseModel, Field
+
+
+def get_model():
+    """
+    Get a model based on the environment variable.
+    """
+    model = os.getenv("MODEL", "openai")
+
+    if model == "openai":
+        return ChatOpenAI(temperature=0, model="gpt-4o")
+    if model == "anthropic":
+        return ChatAnthropic(
+            temperature=0,
+            model_name="claude-3-5-sonnet-20240620",
+            timeout=None,
+            stop=None
+        )
+
+    raise ValueError("Invalid model specified")
+
 
 class EmailAgentState(MessagesState):
     """Email Agent State"""
     email: str
 
+class EmailTool(BaseModel):
+    """
+    Write an email.
+    """
+    the_email: str = Field(description="The email to be written.")
+
+
 async def email_node(state: EmailAgentState, config: RunnableConfig):
     """
     Write an email.
@@ -25,42 +55,24 @@ async def email_node(state: EmailAgentState, config: RunnableConfig):
         emit_tool_calls=True,
     )
 
-    system_message = "You write emails."
-
-    email_tool = {
-        'name': 'write_email',
-        'description': """Write an email.""",
-        'parameters': {
-            'type': 'object',
-            'properties': {
-                'the_email': {
-                    'description': """The email""",
-                    'type': 'string',                    
-                }
-            },
-            'required': ['the_email']
-        }
-    }
+    instructions = "You write emails."
 
-    email_model = ChatOpenAI(model="gpt-4o").bind_tools(
-        [email_tool],
-        parallel_tool_calls=False,
-        tool_choice="write_email"
+    email_model = get_model().bind_tools(
+        [EmailTool],
+        tool_choice="EmailTool"
     )
 
-    print("GENERATING EMAIL")
     response = await email_model.ainvoke([
         *state["messages"],
-        SystemMessage(
-            content=system_message
+        HumanMessage(
+            content=instructions
         )
     ], config)
 
     tool_calls = cast(Any, response).tool_calls
 
     email = tool_calls[0]["args"]["the_email"]
 
-    
     return {
         "email": email,
     }
@@ -90,22 +102,11 @@ async def send_email_node(state: EmailAgentState, config: RunnableConfig):
         "messages": state["messages"],
     }
 
-def route(state: EmailAgentState):
-    """Route to the appropriate node."""
-
-    print("ROUTING")
-    print(state.get("email", None) is not None)
-    print("---")
-
-    if state.get("email", None) is not None:
-        return "send_email_node"
-    return "email_node"
 
 workflow = StateGraph(EmailAgentState)
 workflow.add_node("email_node", email_node)
 workflow.add_node("send_email_node", send_email_node)
 workflow.set_entry_point("email_node")
-workflow.set_conditional_entry_point(route)
 
 workflow.add_edge("email_node", "send_email_node")
 workflow.add_edge("send_email_node", END)