Update ki_gen/utils.py

ki_gen/utils.py

@@ -7,7 +7,10 @@ from typing import Annotated, Union
 from typing_extensions import TypedDict
 
 from langchain_community.graphs import Neo4jGraph
-from langchain_groq import ChatGroq
+# Remove ChatGroq import
+# from langchain_groq import ChatGroq 
+# Add ChatGoogleGenerativeAI import
+from langchain_google_genai import ChatGoogleGenerativeAI 
 from langchain_openai import ChatOpenAI
 
 from langgraph.checkpoint.sqlite import SqliteSaver
@@ -72,13 +75,17 @@ def _set_env(var: str, value: str = None):
             os.environ[var] = getpass.getpass(f"{var}: ")
 
 
-def init_app(openai_key : str = None, groq_key : str = None, langsmith_key : str = None):
+# Remove groq_key parameter
+def init_app(openai_key : str = None, langsmith_key : str = None): 
     """
     Initialize app with user api keys and sets up proxy settings
     """
-    _set_env("GROQ_API_KEY", value=os.getenv("groq_api_key"))
+    # Remove setting GROQ_API_KEY
+    # _set_env("GROQ_API_KEY", value=os.getenv("groq_api_key")) 
     _set_env("LANGSMITH_API_KEY", value=os.getenv("langsmith_api_key"))
     _set_env("OPENAI_API_KEY", value=os.getenv("openai_api_key"))
+    # Make sure GEMINI_API_KEY is set if needed elsewhere, though ChatGoogleGenerativeAI reads it automatically
+    _set_env("GEMINI_API_KEY", value=os.getenv("gemini_api_key"))
     os.environ["LANGSMITH_TRACING_V2"] = "true"
     os.environ["LANGCHAIN_PROJECT"] = "3GPP Test"
     
@@ -93,14 +100,22 @@ def clear_memory(memory, thread_id: str = "") -> None:
     #checkpoint = base.empty_checkpoint()
     #memory.put(config={"configurable": {"thread_id": thread_id}}, checkpoint=checkpoint, metadata={})  
 
-def get_model(model : str = "deepseek-r1-distill-llama-70b"):
+# Update get_model to use ChatGoogleGenerativeAI
+def get_model(model : str = "gemini-2.0-flash"): 
     """
     Wrapper to return the correct llm object depending on the 'model' param
     """
     if model == "gpt-4o":
         llm = ChatOpenAI(model=model, base_url="https://llm.synapse.thalescloud.io/")
+    # Check for gemini models
+    elif model.startswith("gemini"):
+        # Pass the API key explicitly, although it often reads from env var by default
+        llm = ChatGoogleGenerativeAI(model=model, google_api_key=os.getenv("gemini_api_key")) 
     else:
-        llm = ChatGroq(model=model)
+        # Fallback or handle other models if necessary, maybe raise an error
+        # For now, defaulting to Gemini if model name doesn't match others
+        print(f"Warning: Model '{model}' not explicitly handled. Defaulting to Gemini.")
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", google_api_key=os.getenv("gemini_api_key"))
     return llm
 
 
@@ -148,5 +163,4 @@ class DocProcessorState(TypedDict):
     valid_docs : list[Union[str, dict]]
     docs_in_processing : list
     process_steps : list[Union[str,dict]]
-    current_process_step : int
-    
+    current_process_step : int