Update ki_gen/planner.py

ki_gen/planner.py

@@ -4,7 +4,12 @@ import re
 from typing import Annotated
 from typing_extensions import TypedDict
 
-from langchain_groq import ChatGroq
+# Remove ChatGroq import
+# from langchain_groq import ChatGroq
+# Add ChatGoogleGenerativeAI import
+from langchain_google_genai import ChatGoogleGenerativeAI
+import os # Add os import
+
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_community.graphs import Neo4jGraph
@@ -15,11 +20,11 @@ from langgraph.graph import add_messages
 from ki_gen.prompts import PLAN_GEN_PROMPT, PLAN_MODIFICATION_PROMPT
 from ki_gen.data_retriever import build_data_retriever_graph
 from ki_gen.data_processor import build_data_processor_graph
-from ki_gen.utils import ConfigSchema, State, HumanValidationState, DocProcessorState, DocRetrieverState
+# Import get_model which now handles Gemini
+from ki_gen.utils import ConfigSchema, State, HumanValidationState, DocProcessorState, DocRetrieverState, get_model 
 from langgraph.checkpoint.sqlite import SqliteSaver
 
 
-
 ##########################################################################
 ######                       NODES DEFINITION                       ######
 ##########################################################################
@@ -33,108 +38,138 @@ If needed, give me an updated plan to follow this instruction. If your plan alre
     output = HumanMessage(content=prompt)
     return {"messages" : [output]}
 
-
-def error_chatbot_groq(error, model_name, query): # Pass model_name instead of llm_groq object
-    # Switch API key logic...
-    if os.environ["GROQ_API_KEY"] == os.getenv("groq_api_key"):
-        os.environ["GROQ_API_KEY"] = os.getenv("groq_api_key2")
-    elif os.environ["GROQ_API_KEY"] == os.getenv("groq_api_key2"):
-        os.environ["GROQ_API_KEY"] = os.getenv("groq_api_key3")
-    else:
-        os.environ["GROQ_API_KEY"] = os.getenv("groq_api_key")
-
-    # Re-initialize the model *after* switching the key
-    try:
-        # Use the model_name passed in
-        llm_groq_retry = ChatGroq(model=model_name)
-        # Pass the original query messages
-        return {"messages" : [llm_groq_retry.invoke(query)]}
-    except Exception as retry_error:
-        # Handle potential error during retry
-        print(f"Error during retry: {retry_error}")
-        # Decide what to return or raise here
-        return {"messages": [SystemMessage(content=f"Failed to process after retry: {retry_error}")]}
-
+# Remove Groq-specific error handler
+# def error_chatbot_groq(error, model_name, query): ...
 
 # Wrappers to call LLMs on the state messsages field
-def chatbot_llama(state: State):
+# Simplify: Use get_model directly or a single chatbot function
+def chatbot_node(state: State, config: ConfigSchema):
+    """Generic chatbot node using the main_llm from config."""
+    model_name = config["configurable"].get("main_llm") or "gemini-2.0-flash"
+    llm = get_model(model_name)
     try:
-        llm_llama = ChatGroq(model="llama3-70b-8192")
-        return {"messages" : [llm_llama.invoke(state["messages"])]}
-    except Exception as error:
-        error_chatbot_groq(error,llm_llama,state["messages"])
-def chatbot_mixtral(state: State):
-    print(state)
-    llm_mixtral = ChatGroq(model="deepseek-r1-distill-llama-70b")
-    print(llm_mixtral)
-    return {"messages" : [llm_mixtral.invoke(state["messages"])]}
-    # except Exception as error:
-        # error_chatbot_groq(error,llm_mixtral,state["messages"])
-def chatbot_openai(state: State):
-    llm_openai = ChatOpenAI(model='gpt-4o', base_url="https://llm.synapse.thalescloud.io/")
-    return {"messages" : [llm_openai.invoke(state["messages"])]}
-
-chatbots = {"gpt-4o" : chatbot_openai,
-            "deepseek-r1-distill-llama-70b" : chatbot_mixtral,
-            "llama3-70b-8192" : chatbot_llama
-            }
+        # Check if messages exist and are not empty
+        if "messages" in state and state["messages"]:
+             response = llm.invoke(state["messages"])
+             return {"messages": [response]}
+        else:
+             print("Warning: No messages found in state for chatbot_node.")
+             # Return state unchanged or an empty message list?
+             return {} # Or {"messages": []}
+    except Exception as e:
+        print(f"Error invoking model {model_name}: {e}")
+        # Handle error, maybe return an error message or empty dict
+        return {"messages": [SystemMessage(content=f"Error during generation: {e}")]}
+
 
+# Remove old chatbot functions (chatbot_llama, chatbot_mixtral, chatbot_openai)
+# Replace the chatbots dictionary with direct calls to the generic function or specific models via get_model
+# This simplifies planner.py, relying on utils.py and config for model selection.
 
 def parse_plan(state: State):
     """
     This node parses the generated plan and writes in the 'store_plan' field of the state
     """
-    plan = state["messages"][-3].content
-    store_plan = re.split("\d\.", plan.split("Plan:\n")[1])[1:]
+    # Find the AI message likely containing the plan (often the second to last if validate_node was used)
+    plan_message_content = ""
+    if "messages" in state and len(state["messages"]) >= 1:
+         # Search backwards for the plan, as its position might vary
+         for msg in reversed(state["messages"]):
+              if hasattr(msg, 'content') and "Plan:" in msg.content and "<END_OF_PLAN>" in msg.content:
+                   plan_message_content = msg.content
+                   break # Found the plan
+    
+    if not plan_message_content:
+         print("Error: Could not find plan message in state.")
+         # Handle error: maybe return current state or raise an exception
+         return state # Return unchanged state if plan not found
+
+    store_plan = []
     try:
-        store_plan[len(store_plan) - 1] = store_plan[len(store_plan) - 1].split("<END_OF_PLAN>")[0]
+        # Improved parsing: handle potential variations in formatting
+        plan_section = plan_message_content.split("Plan:")[1].split("<END_OF_PLAN>")[0]
+        # Split by numbered steps, removing empty entries
+        store_plan = [step.strip() for step in re.split(r"\n\s*\d+\.\s*", plan_section) if step.strip()]
     except Exception as e:
-        print(f"Error while removing <END_OF_PLAN> : {e}")
+        print(f"Error while parsing plan: {e}")
+        # Handle parsing error, potentially keep store_plan empty or log the error
+        store_plan = [] # Reset plan on error
 
     return {"store_plan" : store_plan}
 
+# Update get_detailed_query to use get_model and default model
+def get_detailed_query(context : list, model : str = "gemini-2.0-flash"):
+    """
+    Simple helper function for the detail_step node
+    """
+    llm = get_model(model) # Use get_model
+    try:
+        return llm.invoke(context)
+    except Exception as e:
+        print(f"Error in get_detailed_query with model {model}: {e}")
+        # Return a default message or raise error
+        return SystemMessage(content=f"Error generating detailed query: {e}")
+
+
 def detail_step(state: State, config: ConfigSchema):
     """
     This node updates the value of the 'current_plan_step' field and defines the query to be used for the data_retriever.
     """
-    print("test")
-    print(state)
+    print("Entering detail_step") # Debug print
+    print(f"Current state keys: {state.keys()}") # Debug print
+
+    # Initialize current_plan_step if not present
+    current_plan_step = state.get("current_plan_step", -1) + 1 
+    
+    # Ensure store_plan exists and has enough steps
+    store_plan = state.get("store_plan", [])
+    if not store_plan or current_plan_step >= len(store_plan):
+        print(f"Warning: Plan step {current_plan_step} out of bounds or plan is empty.")
+        # Decide how to handle: end graph, return error state?
+        # For now, let's prevent index error and maybe signal an issue
+        # Returning an empty query might halt progress or cause issues downstream
+        return {"current_plan_step": current_plan_step, 'query' : "Error: Plan step unavailable.", "valid_docs" : []}
 
-    if 'current_plan_step' in state.keys():
-        print("all good chief")
-    else:
-        state["current_plan_step"] = None
 
-    current_plan_step = state["current_plan_step"] + 1 if state["current_plan_step"] is not None else 0 # We just began a new step so we will increase current_plan_step at the end
+    plan_step_description = store_plan[current_plan_step]
+    
     if config["configurable"].get("use_detailed_query"):
         prompt = HumanMessage(f"""Specify what additional information you need to proceed with the next step of your plan :
-                            Step {current_plan_step + 1} : {state['store_plan'][current_plan_step]}""")
-        query = get_detailed_query(context = state["messages"] + [prompt], model=config["configurable"].get("main_llm"))
-        return {"messages" : [prompt, query], "current_plan_step": current_plan_step, 'query' : query}
-    
-    return {"current_plan_step": current_plan_step, 'query' : state["store_plan"][current_plan_step], "valid_docs" : []}
+                            Step {current_plan_step + 1} : {plan_step_description}""")
+        # Ensure messages exist before appending
+        current_messages = state.get("messages", [])
+        query_message = get_detailed_query(context = current_messages + [prompt], model=config["configurable"].get("main_llm", "gemini-2.0-flash"))
+        query_content = query_message.content if hasattr(query_message, 'content') else "Error: Could not get detailed query content."
+        return {"messages" : [prompt, query_message], "current_plan_step": current_plan_step, 'query' : query_content, "valid_docs": state.get("valid_docs", [])} # Ensure valid_docs is preserved
+
+    # If not using detailed query, use the plan step description directly
+    return {"current_plan_step": current_plan_step, 'query' : plan_step_description, "valid_docs" : state.get("valid_docs", [])} # Ensure valid_docs is preserved
 
-def get_detailed_query(context : list, model : str = "deepseek-r1-distill-llama-70b"):
-    """
-    Simple helper function for the detail_step node
-    """
-    if model == 'gpt-4o':
-        llm = ChatOpenAI(model=model, base_url="https://llm.synapse.thalescloud.io/")
-    else:
-        llm = ChatGroq(model=model)
-    return llm.invoke(context)
 
 def concatenate_data(state: State):
     """
     This node concatenates all the data that was processed by the data_processor and inserts it in the state's messages
     """
+     # Ensure valid_docs exists and current_plan_step is valid
+    valid_docs_content = state.get("valid_docs", "No processed documents available.")
+    current_plan_step = state.get("current_plan_step", -1)
+    store_plan = state.get("store_plan", [])
+
+    if current_plan_step < 0 or current_plan_step >= len(store_plan):
+         print(f"Warning: Invalid current_plan_step ({current_plan_step}) in concatenate_data.")
+         # Handle error - maybe return an error message
+         step_description = "Error: Current plan step invalid."
+    else:
+         step_description = store_plan[current_plan_step]
+
+
     prompt = f"""#########TECHNICAL INFORMATION ############
-{str(state["valid_docs"])}
+{str(valid_docs_content)}
 
 ########END OF TECHNICAL INFORMATION#######
 
-Using the information provided above, proceed with step {state['current_plan_step'] + 1} of your plan : 
-    {state['store_plan'][state['current_plan_step']]}
+Using the information provided above, proceed with step {current_plan_step + 1} of your plan : 
+    {step_description}
 """
     
     return {"messages": [HumanMessage(content=prompt)]}
@@ -142,18 +177,32 @@ Using the information provided above, proceed with step {state['current_plan_ste
 
 def human_validation(state: HumanValidationState) -> HumanValidationState:
     """
-    Dummy node to interrupt before 
+    Dummy node to interrupt before processing, can be used for manual validation later.
     """
-    return {'process_steps' : []}
+    # Defaulting to no processing steps needed unless specified elsewhere
+    return {'process_steps' : state.get('process_steps', [])} 
+
 
 def generate_ki(state: State):
     """
     This node inserts the prompt to begin Key Issues generation
     """
-    print(f"THIS IS THE STATE FOR CURRENT PLAN STEP IN GENERATE_KI : {state}")
+    print(f"THIS IS THE STATE FOR CURRENT PLAN STEP IN GENERATE_KI : {state.get('current_plan_step')}")
+
+    current_plan_step = state.get("current_plan_step", -1)
+    store_plan = state.get("store_plan", [])
+
+    # Check if the next step exists in the plan
+    next_step_index = current_plan_step + 1
+    if next_step_index < 0 or next_step_index >= len(store_plan):
+        print(f"Warning: Invalid next plan step ({next_step_index}) for KI generation.")
+        step_description = "Error: Plan step for KI generation unavailable."
+    else:
+        step_description = store_plan[next_step_index]
 
-    prompt = f"""Using the information provided above, proceed with step 4 of your plan to provide the user with NEW and INNOVATIVE Key Issues :
-    {state['store_plan'][state['current_plan_step'] + 1]}"""
+
+    prompt = f"""Using the information provided above, proceed with step {next_step_index + 1} of your plan to provide the user with NEW and INNOVATIVE Key Issues :
+    {step_description}"""
 
     return {"messages" : [HumanMessage(content=prompt)]}
 
@@ -161,8 +210,19 @@ def detail_ki(state: State):
     """
     This node inserts the last prompt to detail the generated Key Issues
     """
-    prompt = f"""Using the information provided above, proceed with step 5 of your plan to provide the user with NEW and INNOVATIVE Key Issues :
-    {state['store_plan'][state['current_plan_step'] + 2]}"""
+    current_plan_step = state.get("current_plan_step", -1)
+    store_plan = state.get("store_plan", [])
+
+    # Check if the step after next exists in the plan
+    detail_step_index = current_plan_step + 2
+    if detail_step_index < 0 or detail_step_index >= len(store_plan):
+        print(f"Warning: Invalid plan step ({detail_step_index}) for KI detailing.")
+        step_description = "Error: Plan step for KI detailing unavailable."
+    else:
+        step_description = store_plan[detail_step_index]
+
+    prompt = f"""Using the information provided above, proceed with step {detail_step_index + 1} of your plan to provide the user with NEW and INNOVATIVE Key Issues :
+    {step_description}"""
 
     return {"messages" : [HumanMessage(content=prompt)]}
 
@@ -174,44 +234,54 @@ def validate_plan(state: State):
     """
     Whether to regenerate the plan or to parse it
     """
-    if "messages" in state and "My plan is correct" in state["messages"][-1].content:
-        return "parse"
+    # Check the last message for "My plan is correct"
+    if "messages" in state and state["messages"]:
+        last_message = state["messages"][-1]
+        if hasattr(last_message, 'content') and "My plan is correct" in last_message.content:
+             return "parse"
+    # Default to validate (regenerate) if condition not met or messages are missing
     return "validate"
 
 def next_plan_step(state: State, config: ConfigSchema):
     """
     Proceed to next plan step (either generate KI or retrieve more data)
     """
-    if (state["current_plan_step"] == 2) and (config["configurable"].get('plan_method') == "modification"):
-        return "generate_key_issues"
-    if state["current_plan_step"] == len(state["store_plan"]) - 1:
+    current_plan_step = state.get("current_plan_step", -1)
+    store_plan_len = len(state.get("store_plan", []))
+
+    # Simplified logic: go to KI generation if it's the last step based on plan length
+    if current_plan_step >= store_plan_len - 1:
         return "generate_key_issues"
     else:
         return "detail_step"
+
     
 def detail_or_data_retriever(state: State, config: ConfigSchema):
     """
-    Detail the query to use for data retrieval or not
+    Decide whether to detail the query or go straight to data retrieval.
     """
+    # Check configuration if detailed query is needed
     if config["configurable"].get("use_detailed_query"):
-        return "chatbot_detail"
+         # Need to invoke the LLM to get the detailed query
+        return "chatbot_detail" 
     else:
+        # Use the plan step directly as the query
         return "data_retriever"
     
 def retrieve_or_process(state: State):
     """
-    Process the retrieved docs or keep retrieving
+    Process the retrieved docs or keep retrieving (based on human_validated flag).
     """
-    if state['human_validated']:
+    # Check the 'human_validated' flag in the state
+    # This flag needs to be set externally (e.g., by Streamlit UI or another mechanism)
+    # before this node is reached after data_retriever.
+    if state.get('human_validated'): 
         return "process"
-    return "retrieve"
-    # while True:
-    #     user_input = input(f"{len(state['valid_docs'])} were retreived. Do you want more documents (y/[n]) : ")
-    #     if user_input.lower() == "y":
-    #         return "retrieve"
-    #     if not user_input or user_input.lower() == "n":
-    #         return "process"
-    #     print("Please answer with 'y' or 'n'.\n")
+    else:
+        # If not validated, loop back to retrieve more (or wait for validation)
+        # This assumes data_retriever might be called again or the graph waits.
+        # In the Streamlit app, the human_validation node allows setting this flag.
+        return "retrieve" 
 
 
 def build_planner_graph(memory, config):
@@ -222,30 +292,36 @@ def build_planner_graph(memory, config):
 
     graph_doc_retriever = build_data_retriever_graph(memory)
     graph_doc_processor = build_data_processor_graph(memory)
-    graph_builder.add_node("chatbot_planner", chatbots[config["main_llm"]])
+
+    # Use the generic chatbot node function
+    graph_builder.add_node("chatbot_planner", lambda state: chatbot_node(state, config))
     graph_builder.add_node("validate", validate_node)
-    graph_builder.add_node("chatbot_detail", chatbot_llama)
+     # Add node for chatbot interaction when detailed query is needed
+    graph_builder.add_node("chatbot_detail", lambda state: chatbot_node(state, config))
     graph_builder.add_node("parse", parse_plan)
-    graph_builder.add_node("detail_step", detail_step)
-    graph_builder.add_node("data_retriever", graph_doc_retriever, input=DocRetrieverState)
-    graph_builder.add_node("human_validation", human_validation)
-    graph_builder.add_node("data_processor", graph_doc_processor, input=DocProcessorState)
+    # Pass config to detail_step as it needs it now
+    graph_builder.add_node("detail_step", lambda state: detail_step(state, config)) 
+    graph_builder.add_node("data_retriever", graph_doc_retriever) # Input mapping happens automatically if state keys match
+    graph_builder.add_node("human_validation", human_validation) # Needs input mapping if HumanValidationState differs significantly
+    graph_builder.add_node("data_processor", graph_doc_processor) # Needs input mapping if DocProcessorState differs significantly
     graph_builder.add_node("concatenate_data", concatenate_data)
-    graph_builder.add_node("chatbot_exec_step", chatbots[config["main_llm"]])
+    # Use the generic chatbot node function
+    graph_builder.add_node("chatbot_exec_step", lambda state: chatbot_node(state, config))
     graph_builder.add_node("generate_ki", generate_ki)
-    graph_builder.add_node("chatbot_ki", chatbots[config["main_llm"]])
+    # Use the generic chatbot node function
+    graph_builder.add_node("chatbot_ki", lambda state: chatbot_node(state, config))
     graph_builder.add_node("detail_ki", detail_ki)
-    graph_builder.add_node("chatbot_final", chatbots[config["main_llm"]])
+    # Use the generic chatbot node function
+    graph_builder.add_node("chatbot_final", lambda state: chatbot_node(state, config))
 
+    # Define edges
     graph_builder.add_edge("validate", "chatbot_planner")
     graph_builder.add_edge("parse", "detail_step")
 
+    # Edge from chatbot_detail (after getting detailed query) to data_retriever
+    graph_builder.add_edge("chatbot_detail", "data_retriever") 
 
-    # graph_builder.add_edge("detail_step", "chatbot2")
-    graph_builder.add_edge("chatbot_detail", "data_retriever")
     graph_builder.add_edge("data_retriever", "human_validation")
-
-
     graph_builder.add_edge("data_processor", "concatenate_data")
     graph_builder.add_edge("concatenate_data", "chatbot_exec_step")
     graph_builder.add_edge("generate_ki", "chatbot_ki")
@@ -253,15 +329,18 @@ def build_planner_graph(memory, config):
     graph_builder.add_edge("detail_ki", "chatbot_final")
     graph_builder.add_edge("chatbot_final", "__end__")
 
+    # Define conditional edges
     graph_builder.add_conditional_edges(
         "detail_step",
-        detail_or_data_retriever,
+        # Pass config to the conditional function
+        lambda state: detail_or_data_retriever(state, config), 
         {"chatbot_detail": "chatbot_detail", "data_retriever": "data_retriever"}
     )
     graph_builder.add_conditional_edges(
         "human_validation",
         retrieve_or_process,
-        {"retrieve" : "data_retriever", "process" : "data_processor"}
+        # Map 'retrieve' back to 'data_retriever' node, 'process' to 'data_processor'
+        {"retrieve" : "data_retriever", "process" : "data_processor"} 
     )
     graph_builder.add_conditional_edges(
         "chatbot_planner",
@@ -270,13 +349,18 @@ def build_planner_graph(memory, config):
     )
     graph_builder.add_conditional_edges(
         "chatbot_exec_step",
-        next_plan_step,
+         # Pass config to the conditional function
+        lambda state: next_plan_step(state, config),
         {"generate_key_issues" : "generate_ki", "detail_step": "detail_step"}
     )
 
-    graph_builder.set_entry_point("chatbot_planner")
+    # Set entry point
+    graph_builder.set_entry_point("chatbot_planner") 
+    
+    # Compile the graph
     graph = graph_builder.compile(
         checkpointer=memory,
-        interrupt_after=["parse", "chatbot_exec_step", "chatbot_final", "data_retriever"],
+        # Define interrupt points if needed for human interaction or debugging
+        interrupt_after=["human_validation", "chatbot_final"], 
     )
     return graph