Move trigger_ai_message_with_tool_call to conversation/generate.py

app.py

@@ -5,8 +5,6 @@ from langchain.chat_models import init_chat_model
 from langchain_core.tools import tool
 from langchain_aws import BedrockEmbeddings
 from langchain_qdrant import QdrantVectorStore
-from langchain.schema import AIMessage, HumanMessage
-from langchain_core.messages.tool import ToolCall
 
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import MessagesState, StateGraph, END
@@ -19,7 +17,7 @@ from qdrant_client.http.models import Distance, VectorParams, SparseVectorParams
 
 import logging_config as _
 # from conversation.main import graph
-from conversation.generate import generate
+from conversation.generate import generate, trigger_ai_message_with_tool_call
 from conversation.source_history import prettify_source_history, build_source_history_object
 from ingestion.main import ingest_document
 from tools.langfuse_client import get_langfuse_handler
@@ -73,40 +71,6 @@ def retrieve(query: str):
     )
     return serialized, retrieved_docs
 
-def trigger_ai_message_with_tool_call(state: MessagesState) -> AIMessage:
-    """
-    Takes the last user message from the state and returns an AIMessage
-    with example tool_calls populated.
-    
-    Args:
-        state (dict): A dictionary with a 'messages' key containing a list of LangChain messages.
-
-    Returns:
-        AIMessage: An AIMessage with tool_calls based on the last user message.
-    """
-    
-    # Filter for user messages
-    user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
-    
-    if not user_messages:
-        raise ValueError("No user messages found in the previous messages.")
-    
-    last_user_msg = user_messages[-1]
-
-    tool_call = ToolCall(
-        name="retrieve",
-        args={"query": last_user_msg.content},
-        id="tool_call_1"
-    )
-
-    # Construct the AIMessage with tool_calls
-    ai_message = AIMessage(
-        content="Calling the retrieve function...",
-        tool_calls=[tool_call]
-    )
-
-    return {"messages": [ai_message]}
-
 
 graph_builder = StateGraph(MessagesState)
 memory = MemorySaver()

conversation/generate.py

@@ -1,6 +1,7 @@
 import structlog
 from langchain.chat_models import init_chat_model
 from langchain_core.messages import SystemMessage, AIMessage, HumanMessage
+from langchain_core.messages.tool import ToolCall
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import (
     ChatPromptTemplate,
@@ -132,4 +133,38 @@ def generate(state: MessagesState):
         "messages": main_answer,
         "llm-answer": structured_response.answer,
         "sources": citations
-    }
+    }
+
+def trigger_ai_message_with_tool_call(state: MessagesState) -> AIMessage:
+    """
+    Takes the last user message from the state and returns an AIMessage
+    with example tool_calls populated.
+    
+    Args:
+        state (dict): A dictionary with a 'messages' key containing a list of LangChain messages.
+
+    Returns:
+        AIMessage: An AIMessage with tool_calls based on the last user message.
+    """
+    
+    # Filter for user messages
+    user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
+    
+    if not user_messages:
+        raise ValueError("No user messages found in the previous messages.")
+    
+    last_user_msg = user_messages[-1]
+
+    tool_call = ToolCall(
+        name="retrieve",
+        args={"query": last_user_msg.content},
+        id="tool_call_1"
+    )
+
+    # Construct the AIMessage with tool_calls
+    ai_message = AIMessage(
+        content="Calling the retrieve function...",
+        tool_calls=[tool_call]
+    )
+
+    return {"messages": [ai_message]}