Update ki_gen/data_processor.py

ki_gen/data_processor.py

@@ -4,22 +4,21 @@
 from langchain_openai import ChatOpenAI
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-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 for getenv
+
 from langgraph.graph import StateGraph
 from llmlingua import PromptCompressor
 
-from ki_gen.utils import ConfigSchema, DocProcessorState, get_model, format_doc
+# Import get_model which now handles Gemini
+from ki_gen.utils import ConfigSchema, DocProcessorState, get_model, format_doc 
 from langgraph.checkpoint.sqlite import SqliteSaver
 
 
-
-
-# compressed_prompt = llm_lingua.compress_prompt(prompt, instruction="", question="", target_token=200)
-
-## Or use the quantation model, like TheBloke/Llama-2-7b-Chat-GPTQ, only need <8GB GPU memory.
-## Before that, you need to pip install optimum auto-gptq
-# llm_lingua = PromptCompressor("TheBloke/Llama-2-7b-Chat-GPTQ", model_config={"revision": "main"})
-
+# ... (rest of the imports and llm_lingua functions remain the same)
 
 
 # Requires ~2GB of RAM
@@ -61,6 +60,7 @@ def compress(state: DocProcessorState, config: ConfigSchema):
     
     return {"docs_in_processing": doc_process_histories, "current_process_step" : state["current_process_step"] + 1}
 
+# Update default model
 def summarize_docs(state: DocProcessorState, config: ConfigSchema):
     """
     This node summarizes all docs in state["valid_docs"]
@@ -74,12 +74,11 @@ Document:
     sysmsg = ChatPromptTemplate.from_messages([
         ("system", prompt)
     ])
-    model = config["configurable"].get("summarize_model") or "deepseek-r1-distill-llama-70b"
+    # Update default model name
+    model = config["configurable"].get("summarize_model") or "gemini-2.0-flash" 
     doc_process_histories = state["docs_in_processing"]
-    if model == "gpt-4o":
-        llm_summarize = ChatOpenAI(model='gpt-4o', base_url="https://llm.synapse.thalescloud.io/")
-    else:
-        llm_summarize = ChatGroq(model=model)
+    # Use get_model to handle instantiation
+    llm_summarize = get_model(model) 
     summarize_chain = sysmsg | llm_summarize | StrOutputParser()
 
     for doc_process_history in doc_process_histories:
@@ -87,6 +86,7 @@ Document:
 
     return {"docs_in_processing": doc_process_histories, "current_process_step": state["current_process_step"] + 1}
 
+# Update default model
 def custom_process(state: DocProcessorState):
     """
     Custom processing step, params are stored in a dict in state["process_steps"][state["current_process_step"]]
@@ -96,13 +96,15 @@ def custom_process(state: DocProcessorState):
     """
 
     processing_params = state["process_steps"][state["current_process_step"]]
-    model = processing_params.get("processing_model") or "deepseek-r1-distill-llama-70b"
+    # Update default model name
+    model = processing_params.get("processing_model") or "gemini-2.0-flash" 
     user_prompt = processing_params["prompt"]
     context = processing_params.get("context") or [0]
     doc_process_histories = state["docs_in_processing"]
     if not isinstance(context, list):
         context = [context]
 
+    # Use get_model
     processing_chain = get_model(model=model) | StrOutputParser()
 
     for doc_process_history in doc_process_histories:
@@ -113,6 +115,8 @@ def custom_process(state: DocProcessorState):
 
     return {"docs_in_processing" : doc_process_histories, "current_process_step" : state["current_process_step"] + 1}
 
+# ... (rest of the file remains the same)
+
 def final(state: DocProcessorState):
     """
     A node to store the final results of processing in the 'valid_docs' field