Update agent.py

agent.py

@@ -464,12 +464,24 @@ hf_token = os.getenv("HF_TOKEN")
 
 
 # Initialize the desired model and parameters
-llm = LiteLLMModel(
-    model_id="gemini/gemini-2.0-flash", 
-    api_key=os.getenv("GEMINI_API_KEY"),
-    max_tokens=8192
+model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Create a text generation pipeline
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=512,
+    temperature=0.7,
+    top_p=0.95,
+    repetition_penalty=1.15
 )
 
+# Create LangChain LLM wrapper
+llm = HuggingFacePipeline(pipeline=pipe)
+
 
 
 # Initialize the LangChain agent with the tool(s) and the model