Update app.py

app.py

@@ -74,26 +74,30 @@ class SpyAgent(BasicAgent):
             raise NotImplementedError
 
     def llm_caller(self, prompt):
-        client = OpenAI(
-            api_key=os.getenv('API_KEY'),
-            base_url=os.getenv('BASE_URL')
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": prompt}
+        ]
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
         )
-        completion = client.chat.completions.create(
-            model=self.model_name,
-            messages=[
-                {'role': 'system', 'content': 'You are a helpful assistant.'},
-                {'role': 'user', 'content': prompt}
-            ],
-            temperature=0
+        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.device)
+
+        generated_ids = self.model.generate(
+            model_inputs.input_ids,
+            max_new_tokens=512
         )
-        try:
-            return completion.choices[0].message.content
-        except Exception as e:
-            print(e)
-            return None
+        generated_ids = [
+            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
+
+        response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        return response
 
 
 if __name__ == '__main__':
     name = 'spy'
-    agent_builder = AgentBuilder(name, agent=SpyAgent(name, model_name=os.getenv('MODEL_NAME')))
+    agent_builder = AgentBuilder(name, agent=SpyAgent(name, model_name='qwen3-coder-plus'))
     agent_builder.start()