Using "history" with ChatPromptTemplate giving Key error

[A-Sur590]

I am writing a agent code that takes user questions, forms SQL statement and execute the SQL and the get the result. When trying to add persistent memory capability using "history", I am getting error from prompt.format_messages(**kwargs). Here's my full code:

class SLAgent:
    def __init__(self):
        self.db_manager = DatabaseManager()
        self.llm_manager = LLMManager()


    def parse_question(self, state: dict, memory) -> dict:
        
        question = state["question"]
        schema = state["table_schema"]

        #Memory testing
        history = memory.load_memory_variables({})["history"]

        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    """You are a data analyst that can help summarize SQL tables and parse user questions about a database. 
                    
                    """,
                ),
            MessagesPlaceholder("history"),
                (
                    "human",
                    "===Database schema:\n{schema}\n\n===User question:\n{question}\n\nIdentify relevant tables and columns:",
                ),
            ]
        )

        output_parser = JsonOutputParser()

        response = self.llm_manager.invoke(prompt, schema=schema, question=question)
        parsed_response = output_parser.parse(response)
        memory.save_context(question,parsed_response)
        return {"parsed_question": parsed_response}

class LLMManager:
    def __init__(self):
        self.llm = ChatOpenAI(
            model_name=app_config.openai_llm_config["model"],
            api_key=app_config.openai_llm_config["api_key"],
            temperature=app_config.openai_llm_config["temperature"],
        )

    def get_llm(self):
        return self.llm
    
    def invoke(self, prompt: ChatPromptTemplate, **kwargs) -> str:
        messages = prompt.format_messages(**kwargs)
        response = self.llm.invoke(messages)
        return response.content

Here's how invoke parse_question:

class WorkflowManager:
    def __init__(self):
        self.sql_agent = SLAgent()
       
        self.memory = ConversationBufferMemory(memory_key="history")
        self.langfuse_callback = CallbackHandler(
            trace_name="Agent",
            session_id=str(uuid.uuid4()),
            user_id="ID 1"
        )

    def create_workflow(self) -> StateGraph:
       
        workflow = StateGraph(input=InputState, output=OutputState)

        # Add nodes to the graph
       
        workflow.add_node("parse_question", lambda state: self.sql_agent.parse_question(state= state,memory=self.memory))
        
        # Define edges
       
        workflow.add_edge("parse_question", "get_unique_nouns")

        return workflow

When I execute, I get as:

KeyError: 'history' at messages = prompt.format_messages(**kwargs)

[dosubot[bot]]

I found a similar issue related to the KeyError 'history' when using ChatPromptTemplate. To resolve this, ensure that you provide the history argument when calling format_messages. Here's how you can adjust your code:

1. Provide the history argument: Make sure to pass the history parameter when invoking format_messages.

   def invoke(self, prompt: ChatPromptTemplate, **kwargs) -> str:
       history = kwargs.get("history", [])
       messages = prompt.format_messages(history=history, **kwargs)
       response = self.llm.invoke(messages)
       return response.content

2. Ensure history is loaded correctly: Verify that history is being loaded properly from memory.

   history = memory.load_memory_variables({}).get("history", [])

By ensuring that history is provided, you should be able to avoid the KeyError ^[1][2].

To continue talking to Dosu, mention @dosu.

[A-Sur590]

Thanks! How do I pass history from history = memory.load_memory_variables({}).get("history", []) to my prompt? Currently I was using MessagesPlaceholder("history") while creating prompt using ChatPromptTemplate.from_messages(). But that would not load the history that I got from memory.load_memory_varaibles()

[dosubot[bot]]

To pass history from memory.load_memory_variables({}).get("history", []) to your prompt using ChatPromptTemplate.from_messages(), you can use a MessagesPlaceholder to incorporate the chat history. Here's how you can do it:

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# Assuming memory is an instance that has the method load_memory_variables
history = memory.load_memory_variables({}).get("history", [])

# Create the prompt template with a MessagesPlaceholder for history
prompt_template = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a data analyst that can help summarize SQL tables and parse user questions about a database."),
        MessagesPlaceholder("chat_history"),
        ("human", "===Database schema:\n{schema}\n\n===User question:\n{question}\n\nIdentify relevant tables and columns:"),
    ]
)

# Use the template with the history
prompt = prompt_template.format(
    input_variables={
        "chat_history": history,
        "schema": schema,
        "question": question
    }
)

# Now you can use the prompt as needed

This code snippet demonstrates how to pass the history variable to the prompt template using a MessagesPlaceholder ^[1].

To continue talking to Dosu, mention @dosu.

[A-Sur590]

is it chat_history or history? In memory.load_memory_variables({}).get("history", []) its "history" but then in MessagesPlaceholder its chat_history? Also is prompt_template.format() different from prompt.format_messages()?

[dosubot[bot]]

You should use "chat_history" in MessagesPlaceholder to match the key used in memory.load_memory_variables({}).get("history", []). Consistency in naming ensures the correct data is retrieved and formatted.

Regarding prompt_template.format() vs. prompt.format_messages():

• format() generates a formatted string based on the template and provided variables.
• format_messages() generates a list of message objects (e.g., SystemMessage, HumanMessage) based on the template and provided variables.

Choose based on whether you need a string or structured message objects ^[1][2].

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu Here's my updated code. I am getting error raise ValueError(f"One input key expected got {prompt_input_keys}") at line memory.save_context(question,response):

class SLAgent:
    def __init__(self):
        self.db_manager = DatabaseManager()
        self.llm_manager = LLMManager()


    def parse_question(self, state: dict, memory) -> dict:
        """Parse user question and identify relevant tables and columns."""
        question = state["question"]
        schema = state["table_schema"]

        #Memory testing
        history = memory.load_memory_variables({}).get("history", [])

        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    """You are a data analyst that can help summarize SQL tables and parse user questions about a database. 
                    
                    """,
                ),
            MessagesPlaceholder("history"),
                (
                    "human",
                    "===Database schema:\n{schema}\n\n===User question:\n{question}\n\nIdentify relevant tables and columns:",
                ),
            ]
        )
        
        output_parser = JsonOutputParser()


        response = self.llm_manager.invoke(prompt, schema=schema, question=question)
        parsed_response = output_parser.parse(response)
        memory.save_context(question,response)
        return {"parsed_question": parsed_response}

class LLMManager:
    def __init__(self):
        self.llm = ChatOpenAI(
            model_name=app_config.openai_llm_config["model"],
            api_key=app_config.openai_llm_config["api_key"],
            temperature=app_config.openai_llm_config["temperature"],
        )

    def get_llm(self):
        return self.llm
    
    def invoke(self, prompt: ChatPromptTemplate, **kwargs) -> str:
       history = kwargs.get("history", [])
        messages = prompt.format_messages(history=history,**kwargs)
        response = self.llm.invoke(messages)
        return response.content

Also in my def parse_question() I am loading "history" from memory using history = memory.load_memory_variables({}).get("history", []) but the variable is not getting passed to any other function. And then I am trying to extract "history" from kwargs in invoke function.

[A-Sur590]

@dosu I was able to get past the error, but seems like the memory functionality is not workign as intended. Here's my latest code that uses Agent workflow:

class SLAgent:
    def __init__(self):
        self.db_manager = DatabaseManager()
        self.llm_manager = LLMManager()


    def parse_question(self, state: dict, memory) -> dict:
        """Parse user question and identify relevant tables and columns."""
        question = state["question"]
        schema = state["table_schema"]

        #Memory testing
        history = memory.load_memory_variables({}).get("history", [])

        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    """You are a data analyst that can help summarize SQL tables and parse user questions about a database. 
                    
                    """,
                ),
            MessagesPlaceholder("history"),
                (
                    "human",
                    "===Database schema:\n{schema}\n\n===User question:\n{question}\n\nIdentify relevant tables and columns:",
                ),
            ]
        )
        
        output_parser = JsonOutputParser()


        response = self.llm_manager.invoke(prompt, schema=schema, question=question)
        parsed_response = output_parser.parse(response)
        return {"parsed_question": parsed_response}

class LLMManager:
    def __init__(self):
        self.llm = ChatOpenAI(
            model_name=app_config.openai_llm_config["model"],
            api_key=app_config.openai_llm_config["api_key"],
            temperature=app_config.openai_llm_config["temperature"],
        )

    def get_llm(self):
        return self.llm
    
    def invoke(self, prompt: ChatPromptTemplate, **kwargs) -> str:
       history = kwargs.get("history", [])
        messages = prompt.format_messages(history=history,**kwargs)
        response = self.llm.invoke(messages)
        return response.content

Here's how the Langgraph workflow is used to invoke:

class WorkflowManager:
    def __init__(self):
        self.sql_agent = SLAgent()
       
        self.memory = ConversationBufferMemory(memory_key="history")
        self.langfuse_callback = CallbackHandler(
            trace_name="Agent",
            session_id=str(uuid.uuid4()),
            user_id="ID 1"
        )

    def create_workflow(self) -> StateGraph:
       
        workflow = StateGraph(input=InputState, output=OutputState)

        # Add nodes to the graph
       
        workflow.add_node("parse_question", lambda state: self.sql_agent.parse_question(state= state,memory=self.memory))
        
        # Define edges
       
        workflow.add_edge("parse_question", "get_unique_nouns")

        return workflow

When I execute this, and do a follow-up conversation, I don't see it picking up from memory at all.

[dosubot[bot]]

To ensure the memory functionality works correctly, you need to make sure that the conversation history is being saved and retrieved properly. Here's how you can adjust your code:

1. Save Context Properly: Ensure that you are saving both the question and the response to the memory.

   memory.save_context({"input": question}, {"output": response})

2. Pass History to the Prompt: When invoking the parse_question method, make sure the history is passed correctly to the invoke method.

3. Update the invoke Method: Ensure that the history is included in the kwargs when calling invoke.

   def parse_question(self, state: dict, memory) -> dict:
       question = state["question"]
       schema = state["table_schema"]
   
       # Load history
       history = memory.load_memory_variables({}).get("history", [])
   
       prompt = ChatPromptTemplate.from_messages(
           [
               ("system", "You are a data analyst that can help summarize SQL tables and parse user questions about a database."),
               MessagesPlaceholder("history"),
               ("human", "===Database schema:\n{schema}\n\n===User question:\n{question}\n\nIdentify relevant tables and columns:"),
           ]
       )
   
       output_parser = JsonOutputParser()
   
       # Pass history to invoke
       response = self.llm_manager.invoke(prompt, schema=schema, question=question, history=history)
       parsed_response = output_parser.parse(response)
   
       # Save context
       memory.save_context({"input": question}, {"output": parsed_response})
   
       return {"parsed_question": parsed_response}

4. Ensure Consistency: Make sure the MessagesPlaceholder key matches the key used in invoke.

By following these steps, your agent should correctly utilize the memory to retain and use conversation history for follow-up questions.

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu It throws error at messages = prompt.format_messages(history=history,**kwargs) :
TypeError: langchain_core.prompts.chat.ChatPromptTemplate.format_messages() got multiple values for keyword argument 'history'

My invoke method is:

def invoke(self, prompt: ChatPromptTemplate, **kwargs) -> str:
       history = kwargs.get("history", [])
        messages = prompt.format_messages(history=history,**kwargs)
        response = self.llm.invoke(messages)
        return response.content

[A-Sur590]

@dosu It throws error at messages = prompt.format_messages(history=history,**kwargs) :
TypeError: langchain_core.prompts.chat.ChatPromptTemplate.format_messages() got multiple values for keyword argument 'history'

My invoke method is:

def invoke(self, prompt: ChatPromptTemplate, **kwargs) -> str:
history = kwargs.get("history", [])
messages = prompt.format_messages(history=history,**kwargs)
response = self.llm.invoke(messages)
return response.content

[dosubot[bot]]

The error occurs because history is being passed twice to format_messages. To fix this, ensure history is not included in **kwargs if it's already being passed explicitly. Here's how you can adjust your invoke method:

def invoke(self, prompt: ChatPromptTemplate, **kwargs) -> str:
    # Remove 'history' from kwargs if it exists
    kwargs.pop("history", None)
    history = kwargs.get("history", [])
    messages = prompt.format_messages(history=history, **kwargs)
    response = self.llm.invoke(messages)
    return response.content

This ensures history is only passed once, resolving the TypeError ^[1][2].

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu it throws following error. The error is at line memory.save_context({"input": question}, {"output": parsed_response}):

memory.save_context({"input": question}, {"output": parsed_response})
File "/usr/local/lib/python3.11/site-packages/langchain/memory/chat_memory.py", line 74, in save_context
[HumanMessage(content=input_str), AIMessage(content=output_str)]
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/langchain_core/messages/ai.py", line 177, in __init__
 super().__init__(content=content, **kwargs)
File "/usr/local/lib/python3.11/site-packages/langchain_core/messages/base.py", line 76, in __init__
super().__init__(content=content, **kwargs)
 File "/usr/local/lib/python3.11/site-packages/langchain_core/load/serializable.py", line 111, in __init__
super().__init__(*args, **kwargs)
File "/usr/local/lib/python3.11/site-packages/pydantic/main.py", line 212, in __init__
validated_self = self.__pydantic_validator__.validate_python(data, self_instance=self)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pydantic_core._pydantic_core.ValidationError: 2 validation errors for AIMessage
content.str
Input should be a valid string [type=string_type, input_value={'is_relevant': True, 're...umns': ['indication']}]}, 

[dosubot[bot]]

Yes, you should handle memory in get_actions and validate_sql if they require context from previous steps. Here's how you can do it:

1. Load and Save Memory: Ensure you load the conversation history at the start and save the context at the end of each function if they depend on previous interactions.

2. Prompt Instructions: You can add instructions in the prompt to handle follow-up questions. This helps the LLM understand when to use previous context.

Here's an example for get_actions and validate_sql:

def get_actions(self, state: dict, memory) -> dict:
    # Load history
    history = memory.load_memory_variables({}).get("history", [])

    # Use history in your logic if needed
    # ...

    # Save context if necessary
    # memory.save_context({"input": some_input}, {"output": some_output})

    return {"unique_actions": list(unique_actions)}

def validate_sql(self, state: dict, memory) -> dict:
    sql_query = state["sql_query"]
    schema = state["schema"]

    # Load history
    history = memory.load_memory_variables({}).get("history", [])

    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are an AI assistant that validates and fixes SQL queries. Consider previous queries if this is a follow-up."),
            MessagesPlaceholder("history"),
            ("human", "===Database schema:\n{schema}\n\n===Generated SQL query:\n{sql_query}"),
        ]
    )

    response = self.llm_manager.invoke(prompt, schema=schema, sql_query=sql_query, history=history)
    result = output_parser.parse(response)

    # Save context
    memory.save_context({"input": sql_query}, {"output": str(result)})

    return {
        "sql_query": result.get("corrected_query", sql_query),
        "sql_valid": result.get("valid", False),
        "sql_issues": result.get("issues", None),
    }

By loading and saving memory in each step, and adding instructions for follow-up handling, you ensure that your agent maintains context across the workflow.

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu Can the same functionality be also achieved using langgraph.checkpoint.memory instead of above method? If yes, can you provide an example how that would work in parse_question and generate_sql along with the WorkflowManager class?

[dosubot[bot]]

Yes, you can achieve similar functionality using langgraph.checkpoint.memory. Here's how you can integrate it into your parse_question, generate_sql, and WorkflowManager classes:

Using langgraph.checkpoint.memory

1. Setup MemorySaver: Use MemorySaver to manage conversation history.

2. Modify WorkflowManager: Integrate MemorySaver into your workflow.

3. Adjust Functions: Use the memory to load and save context.

Example Code

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, MessagesState, StateGraph

class WorkflowManager:
    def __init__(self):
        self.sql_agent = SLAgent()
        self.memory = MemorySaver()

    def create_workflow(self) -> StateGraph:
        workflow = StateGraph(input=InputState, output=OutputState)

        # Add nodes to the graph
        workflow.add_node("parse_question", lambda state: self.sql_agent.parse_question(state=state, memory=self.memory))
        workflow.add_node("get_actions", self.sql_agent.get_actions)
        workflow.add_node("generate_sql", lambda state: self.sql_agent.generate_sql(state=state, memory=self.memory))
        workflow.add_node("validate_sql", self.sql_agent.validate_sql)
        workflow.add_node("execute_sql", self.sql_agent.execute_sql)

        # Define edges
        workflow.add_edge("parse_question", "get_actions")
        workflow.add_edge("get_actions", "generate_sql")
        workflow.add_edge("generate_sql", "validate_sql")
        workflow.add_edge("validate_sql", "execute_sql")

        return workflow

class SLAgent:
    def parse_question(self, state: dict, memory) -> dict:
        question = state["question"]
        schema = state["table_schema"]

        # Load history
        history = memory.load_memory_variables({}).get("history", [])

        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are a data analyst that can help parse user questions about a database."),
                MessagesPlaceholder("history"),
                ("human", "===Database schema:\n{schema}\n\n===User question:\n{question}\n\nIdentify relevant tables and columns:"),
            ]
        )

        output_parser = JsonOutputParser()

        # Pass history to invoke
        response = self.llm_manager.invoke(prompt, schema=schema, question=question, history=history)
        parsed_response = output_parser.parse(response)

        # Save context
        memory.save_context({"input": question}, {"output": str(parsed_response)})

        return {"parsed_question": parsed_response}

    def generate_sql(self, state: dict, memory) -> dict:
        question = state["question"]
        parsed_question = state["parsed_question"]
        unique_actions = state["unique_nouns"]
        schema = state["schema"]

        # Load history
        history = memory.load_memory_variables({}).get("history", [])

        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are an AI assistant that generates SQL queries based on user questions."),
                MessagesPlaceholder("history"),
                ("human", "===Database schema:\n{schema}\n\n===User question:\n{question}\n\n===Relevant tables and columns:\n{parsed_question}\n\n===Unique actions in relevant tables:\n{unique_actions}\n\nGenerate SQL query string"),
            ]
        )

        response = self.llm_manager.invoke(
            prompt,
            schema=schema,
            question=question,
            parsed_question=parsed_question,
            unique_nouns=unique_actions,
            history=history
        )

        # Save context
        memory.save_context({"input": question}, {"output": str(response)})

        return {"sql_query": response}

Key Points

• MemorySaver: Manages the conversation state across sessions.
• Consistent Context Handling: Load and save context in each function that requires it.
• Prompt Instructions: Ensure prompts are clear about using previous context for follow-up questions.

This setup should help maintain context across your workflow using langgraph.checkpoint.memory ^[1][2].

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu if my function get_actions do I save unique_actions in history or inputand output since in this function, no update to question or response happens, only unique_actions are formed/updated?

[dosubot[bot]]

In the get_actions function, since you're forming or updating unique_actions, you should save this specific information in the memory. You can save it as part of the context, even if it's not directly related to the question or response. Here's how you can do it:

def get_actions(self, state: dict, memory) -> dict:
    parsed_question = state["parsed_question"]

    if not parsed_question["is_relevant"]:
        return {"unique_actions": []}

    unique_actions = set()
    for table_info in parsed_question["relevant_tables"]:
        table_name = table_info["table_name"]
        action_columns = table_info["noun_columns"]

        if action_columns:
            column_names = ", ".join(f"{col}" for col in action_columns)
            query = f"SELECT DISTINCT {column_names} FROM {table_name}"
            results = self.db_manager.execute_query(query)
            for row in results:
                unique_actions.update(str(value) for value in row if value)

    # Save unique actions to memory
    memory.save_context({"input": "unique_actions"}, {"output": list(unique_actions)})

    return {"unique_actions": list(unique_actions)}

Key Points

• Save Context: Store unique_actions in the memory to maintain context for future steps.
• Input/Output Structure: Use a descriptive key like "unique_actions" to clarify what is being saved.

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu When I use langgraph.checkpoint.memory import MemorySaver, I am getting error as AttributeError: 'MemorySaver' object has no attribute 'load_memory_variables' at history = memory.load_memory_variables({}).get("history", [])

[dosubot[bot]]

To support both SQL generation and small talk in your workflow, you can introduce a decision-making step that determines whether the input is a SQL-related query or a small talk conversation. Here's how you can adjust your workflow:

Adjusted Workflow

1. Add a Decision Node: Introduce a node that classifies the input as either SQL-related or small talk.

2. Handle Small Talk: Create a separate path in the workflow for handling small talk.

Example Code

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, MessagesState, StateGraph

class WorkflowManager:
    def __init__(self):
        self.sql_agent = SLAgent()
        self.memory = MemorySaver()

    def create_workflow(self) -> StateGraph:
        workflow = StateGraph(input=InputState, output=OutputState)

        # Add nodes to the graph
        workflow.add_node("classify_input", self.classify_input)
        workflow.add_node("parse_question", lambda state: self.sql_agent.parse_question(state=state, memory=self.memory))
        workflow.add_node("get_actions", self.sql_agent.get_actions)
        workflow.add_node("generate_sql", lambda state: self.sql_agent.generate_sql(state=state, memory=self.memory))
        workflow.add_node("validate_sql", self.sql_agent.validate_sql)
        workflow.add_node("execute_sql", self.sql_agent.execute_sql)
        workflow.add_node("handle_small_talk", self.handle_small_talk)

        # Define edges
        workflow.add_edge(START, "classify_input")
        workflow.add_edge("classify_input", "parse_question", condition=lambda state: state["is_sql"])
        workflow.add_edge("parse_question", "get_actions")
        workflow.add_edge("get_actions", "generate_sql")
        workflow.add_edge("generate_sql", "validate_sql")
        workflow.add_edge("validate_sql", "execute_sql")
        workflow.add_edge("classify_input", "handle_small_talk", condition=lambda state: not state["is_sql"])

        return workflow

    def classify_input(self, state: dict) -> dict:
        question = state["question"]
        # Simple logic to classify input; replace with a more sophisticated model if needed
        is_sql = "SELECT" in question.upper() or "FROM" in question.upper()
        return {"is_sql": is_sql}

    def handle_small_talk(self, state: dict) -> dict:
        question = state["question"]
        # Simple response for small talk; replace with a more sophisticated model if needed
        response = f"I'm here to help with your questions! You asked: {question}"
        return {"response": response}

class SLAgent:
    def parse_question(self, state: dict, memory) -> dict:
        question = state["question"]
        schema = state["table_schema"]

        # Load history
        history = memory.load_memory_variables({}).get("history", [])

        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are a data analyst that can help parse user questions about a database."),
                MessagesPlaceholder("history"),
                ("human", "===Database schema:\n{schema}\n\n===User question:\n{question}\n\nIdentify relevant tables and columns:"),
            ]
        )

        output_parser = JsonOutputParser()

        # Pass history to invoke
        response = self.llm_manager.invoke(prompt, schema=schema, question=question, history=history)
        parsed_response = output_parser.parse(response)

        # Save context
        memory.save_context({"input": question}, {"output": str(parsed_response)})

        return {"parsed_question": parsed_response}

    def generate_sql(self, state: dict, memory) -> dict:
        question = state["question"]
        parsed_question = state["parsed_question"]
        unique_actions = state["unique_nouns"]
        schema = state["schema"]

        # Load history
        history = memory.load_memory_variables({}).get("history", [])

        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are an AI assistant that generates SQL queries based on user questions."),
                MessagesPlaceholder("history"),
                ("human", "===Database schema:\n{schema}\n\n===User question:\n{question}\n\n===Relevant tables and columns:\n{parsed_question}\n\n===Unique actions in relevant tables:\n{unique_actions}\n\nGenerate SQL query string"),
            ]
        )

        response = self.llm_manager.invoke(
            prompt,
            schema=schema,
            question=question,
            parsed_question=parsed_question,
            unique_nouns=unique_actions,
            history=history
        )

        # Save context
        memory.save_context({"input": question}, {"output": str(response)})

        return {"sql_query": response}

Key Points

• Classification: Use a simple logic or a model to classify the input as SQL-related or small talk.
• Separate Handling: Create a separate path in the workflow for small talk to ensure it is handled appropriately.
• Flexibility: This setup allows your agent to handle both SQL queries and small talk, enhancing its versatility.

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu I want the LLM to answer to user's small talk and not a static response in handle_small_talk(). The user question is never a SQL statement containing SELECT or FROM. User question can be generic question (like what is sales for product A or what are top products this quarter) which are then transformed into SQL in generate_sql().

[dosubot[bot]]

To enable the LLM to handle small talk dynamically, you can modify the handle_small_talk function to use the LLM for generating responses. Here's how you can adjust your workflow:

Adjusted Workflow for Dynamic Small Talk

1. Use LLM for Small Talk: Instead of a static response, use the LLM to generate responses for small talk.

2. Modify handle_small_talk: Update this function to invoke the LLM with the user's question.

Example Code

class WorkflowManager:
    def __init__(self):
        self.sql_agent = SLAgent()
        self.memory = MemorySaver()

    def create_workflow(self) -> StateGraph:
        workflow = StateGraph(input=InputState, output=OutputState)

        # Add nodes to the graph
        workflow.add_node("classify_input", self.classify_input)
        workflow.add_node("parse_question", lambda state: self.sql_agent.parse_question(state=state, memory=self.memory))
        workflow.add_node("get_actions", self.sql_agent.get_actions)
        workflow.add_node("generate_sql", lambda state: self.sql_agent.generate_sql(state=state, memory=self.memory))
        workflow.add_node("validate_sql", self.sql_agent.validate_sql)
        workflow.add_node("execute_sql", self.sql_agent.execute_sql)
        workflow.add_node("handle_small_talk", lambda state: self.sql_agent.handle_small_talk(state=state, memory=self.memory))

        # Define edges
        workflow.add_edge(START, "classify_input")
        workflow.add_edge("classify_input", "parse_question", condition=lambda state: state["is_sql"])
        workflow.add_edge("parse_question", "get_actions")
        workflow.add_edge("get_actions", "generate_sql")
        workflow.add_edge("generate_sql", "validate_sql")
        workflow.add_edge("validate_sql", "execute_sql")
        workflow.add_edge("classify_input", "handle_small_talk", condition=lambda state: not state["is_sql"])

        return workflow

    def classify_input(self, state: dict) -> dict:
        question = state["question"]
        # Simple logic to classify input; replace with a more sophisticated model if needed
        is_sql = "sales" in question.lower() or "products" in question.lower()
        return {"is_sql": is_sql}

class SLAgent:
    def handle_small_talk(self, state: dict, memory) -> dict:
        question = state["question"]

        # Load history
        history = memory.load_memory_variables({}).get("history", [])

        # Use LLM to generate a response for small talk
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are a friendly assistant ready to chat with the user."),
                MessagesPlaceholder("history"),
                ("human", "{question}"),
            ]
        )

        response = self.llm_manager.invoke(prompt, question=question, history=history)

        # Save context
        memory.save_context({"input": question}, {"output": response})

        return {"response": response}

Key Points

• Dynamic Response: Use the LLM to generate responses for small talk, making interactions more engaging and personalized.
• History Management: Continue to manage conversation history to provide context for the LLM.
• Flexible Classification: Adjust the classification logic to better distinguish between SQL-related queries and small talk.

To continue talking to Dosu, mention @dosu.

[A-Sur590]

@dosu like I said, there is no pattern of keywords that you can use to determine if its a small talk question or if question will form a successful SQL in generate_sql(). Also, the fow should be that if its a small talk, then go to handle_small_talk() and if not, then go to generate_sql()

[dosubot[bot]]

To handle both small talk and SQL-related queries without relying on keyword patterns, you can use the LLM to classify the input dynamically. Here's how you can adjust your workflow:

Adjusted Workflow with LLM Classification

1. Use LLM for Classification: Implement a node that uses the LLM to determine whether the input is small talk or SQL-related.

2. Direct to Appropriate Node: Based on the classification, direct the workflow to either handle_small_talk or generate_sql.

Example Code

class WorkflowManager:
    def __init__(self):
        self.sql_agent = SLAgent()
        self.memory = MemorySaver()

    def create_workflow(self) -> StateGraph:
        workflow = StateGraph(input=InputState, output=OutputState)

        # Add nodes to the graph
        workflow.add_node("classify_input", lambda state: self.sql_agent.classify_input(state=state, memory=self.memory))
        workflow.add_node("generate_sql", lambda state: self.sql_agent.generate_sql(state=state, memory=self.memory))
        workflow.add_node("handle_small_talk", lambda state: self.sql_agent.handle_small_talk(state=state, memory=self.memory))

        # Define edges
        workflow.add_edge(START, "classify_input")
        workflow.add_edge("classify_input", "generate_sql", condition=lambda state: state["is_sql"])
        workflow.add_edge("classify_input", "handle_small_talk", condition=lambda state: not state["is_sql"])

        return workflow

class SLAgent:
    def classify_input(self, state: dict, memory) -> dict:
        question = state["question"]

        # Use LLM to classify the input
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are an assistant that can distinguish between SQL-related queries and small talk."),
                ("human", "Is the following question related to SQL or is it small talk? {question}"),
            ]
        )

        response = self.llm_manager.invoke(prompt, question=question)
        is_sql = "SQL" in response.content  # Adjust based on the LLM's response format

        # Save context
        memory.save_context({"input": question}, {"output": response.content})

        return {"is_sql": is_sql}

    def handle_small_talk(self, state: dict, memory) -> dict:
        question = state["question"]

        # Load history
        history = memory.load_memory_variables({}).get("history", [])

        # Use LLM to generate a response for small talk
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are a friendly assistant ready to chat with the user."),
                MessagesPlaceholder("history"),
                ("human", "{question}"),
            ]
        )

        response = self.llm_manager.invoke(prompt, question=question, history=history)

        # Save context
        memory.save_context({"input": question}, {"output": response})

        return {"response": response}

    def generate_sql(self, state: dict, memory) -> dict:
        question = state["question"]
        parsed_question = state["parsed_question"]
        unique_actions = state["unique_nouns"]
        schema = state["schema"]

        # Load history
        history = memory.load_memory_variables({}).get("history", [])

        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "You are an AI assistant that generates SQL queries based on user questions."),
                MessagesPlaceholder("history"),
                ("human", "===Database schema:\n{schema}\n\n===User question:\n{question}\n\n===Relevant tables and columns:\n{parsed_question}\n\n===Unique actions in relevant tables:\n{unique_actions}\n\nGenerate SQL query string"),
            ]
        )

        response = self.llm_manager.invoke(
            prompt,
            schema=schema,
            question=question,
            parsed_question=parsed_question,
            unique_nouns=unique_actions,
            history=history
        )

        # Save context
        memory.save_context({"input": question}, {"output": str(response)})

        return {"sql_query": response}

Key Points

• LLM Classification: Use the LLM to dynamically classify the input, allowing for more flexible handling of queries.
• Dynamic Routing: Direct the workflow to the appropriate node based on the classification result.
• Context Management: Continue to manage conversation history to provide context for the LLM.

To continue talking to Dosu, mention @dosu.