Update myagent.py

myagent.py

@@ -43,30 +43,46 @@ class BasicAgent:
         
       
     
-# model = OpenAIServerModel(
-#     model_id="gpt-4.1-nano",
-#     api_base="https://api.openai.com/v1",
-#     api_key=os.environ["OPENAI_API_KEY"],
-# )
-
-
-model_id = "bartowski/Llama-3.2-3B-Instruct-GGUF"
-filename = "Llama-3.2-3B-Instruct-Q4_K_M.gguf"
-
-torch_dtype = torch.float32 # could be torch.float16 or torch.bfloat16 too
-tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)
-model_init = AutoModelForCausalLM.from_pretrained(model_id, gguf_file=filename, torch_dtype=torch_dtype)
-
+# Create a wrapper class that matches the expected interface
+class LocalLlamaModel:
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = model.device if hasattr(model, 'device') else 'cpu'
+    
+    def generate(self, prompt: str, max_new_tokens=512, **kwargs):
+        """Generate text using the local model"""
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+        
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                input_ids, 
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=0.7,
+                pad_token_id=self.tokenizer.eos_token_id,
+                **kwargs
+            )
+        
+        # Decode only the new tokens (excluding the input)
+        new_tokens = output_ids[0][input_ids.shape[1]:]
+        output = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+        return output
+    
+    def __call__(self, prompt: str, max_new_tokens=512, **kwargs):
+        """Make the model callable like a function"""
+        return self.generate(prompt, max_new_tokens, **kwargs)
 
-def model(prompt: str, max_new_tokens=512):
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
-    output_ids = model_init.generate(input_ids, max_new_tokens=max_new_tokens)
-    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return output
+# Create the model instance
+model = LocalLlamaModel(model_init, tokenizer)
 
-reviewer_agent= ToolCallingAgent(model=model, tools=[])
-model_agent = ToolCallingAgent(model=model,tools=[fetch_webpage])
-gaia_agent = CodeAgent(tools=[fetch_webpage,get_youtube_title_description,get_youtube_transcript ], model=model)
+# Now create your agents - these should work with the wrapped model
+reviewer_agent = ToolCallingAgent(model=model, tools=[])
+model_agent = ToolCallingAgent(model=model, tools=[fetch_webpage])
+gaia_agent = CodeAgent(
+    tools=[fetch_webpage, get_youtube_title_description, get_youtube_transcript], 
+    model=model
+)
 
 if __name__ == "__main__":
     # Example usage