LLM_as_judge.py - Some fixes for qwen to work - I commit plenty of prints for the future, they will be all removed

examples/cai/agent_patterns/LLM_as_judge.py

@@ -10,11 +10,60 @@ from __future__ import annotations
 
 import asyncio
 import os
+import json
 from dataclasses import dataclass
 from typing import Literal
 from cai.sdk.agents import Agent, ItemHelpers, Runner, TResponseInputItem, OpenAIChatCompletionsModel
 from openai import AsyncOpenAI
+from cai.util import get_ollama_api_base
 
+# Enable debug mode
+#os.environ['CAI_DEBUG'] = '2'
+#os.environ['LITELLM_VERBOSE'] = 'True'
+
+# Force Ollama mode if qwen model is used
+if os.getenv('CAI_MODEL', "qwen2.5:14b").startswith("qwen"): 
+    os.environ['OLLAMA'] = 'true'
+
+# Modify OpenAIChatCompletionsModel._fetch_response_litellm_ollama to debug output
+import cai.sdk.agents.models.openai_chatcompletions
+original_fetch_response_litellm_ollama = cai.sdk.agents.models.openai_chatcompletions.OpenAIChatCompletionsModel._fetch_response_litellm_ollama
+
+async def debug_fetch_response_litellm_ollama(self, kwargs, model_settings, tool_choice, stream, parallel_tool_calls):
+    print("\n[DEBUG] Ollama request parameters:")
+    print(f"Base URL: {get_ollama_api_base().rstrip('/v1')}")
+    print(f"Model name: {kwargs.get('model')}")
+    print(f"Messages: {json.dumps(kwargs.get('messages'))[:200]}...") # Truncated to avoid huge output
+    
+    # Check if the model exists in Ollama
+    import requests
+    try:
+        response = requests.get(f"{get_ollama_api_base().rstrip('/v1')}/api/tags")
+        models = response.json().get("models", [])
+        model_names = [model.get("name") for model in models]
+        print(f"Available Ollama models: {model_names}")
+        
+        model_name = kwargs.get('model')
+        if model_name in model_names:
+            print(f"✅ Model '{model_name}' is available in Ollama")
+        else:
+            print(f"❌ Model '{model_name}' is NOT available in Ollama")
+            similar_models = [name for name in model_names if model_name.split(":")[0] in name]
+            if similar_models:
+                print(f"Similar models available: {similar_models}")
+                
+                # Try with first similar model
+                if similar_models:
+                    print(f"⚠️ Trying with similar model: {similar_models[0]}")
+                    kwargs["model"] = similar_models[0]
+    except Exception as e:
+        print(f"Error checking Ollama models: {e}")
+    
+    # Call the original function
+    return await original_fetch_response_litellm_ollama(self, kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
+
+# Patch the function
+cai.sdk.agents.models.openai_chatcompletions.OpenAIChatCompletionsModel._fetch_response_litellm_ollama = debug_fetch_response_litellm_ollama
 
 # CTF task planner agent (performs planning)
 ctf_task_planner = Agent(
@@ -26,7 +75,7 @@ ctf_task_planner = Agent(
         "Use any feedback to improve your planning."
     ),
     model=OpenAIChatCompletionsModel(
-        model= os.getenv('CAI_MODEL', "qwen2.5:14b"),
+        model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
         openai_client=AsyncOpenAI(),
     ),
     tools=[] 
@@ -50,7 +99,7 @@ ctf_plan_evaluator = Agent[None](
         "Provide actionable feedback. Never approve on the first try."
     ),
     model=OpenAIChatCompletionsModel(
-        model= os.getenv('CAI_MODEL', "qwen2.5:14b"),
+        model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
         openai_client=AsyncOpenAI(),
     ),
     tools=[],  
@@ -65,11 +114,14 @@ async def main() -> None:
     latest_plan: str | None = None
     
     while True:
+        try:
+            print("\n[INFO] Running CTF Task Planner...")
             planner_result = await Runner.run(ctf_task_planner, input_items)
             input_items = planner_result.to_input_list()
             latest_plan = ItemHelpers.text_message_outputs(planner_result.new_items)
             print("CTF task plan generated.")
 
+            print("\n[INFO] Running CTF Plan Evaluator...")
             eval_result = await Runner.run(ctf_plan_evaluator, input_items)
             feedback: EvaluationFeedback = eval_result.final_output
             print(f"Evaluator score: {feedback.score}")
@@ -80,8 +132,16 @@ async def main() -> None:
 
             print("Refining plan based on evaluator feedback...")
             input_items.append({"content": f"Feedback: {feedback.feedback}", "role": "user"})
+        except Exception as e:
+            print(f"Error: {e}")
+            import traceback
+            traceback.print_exc()
+            break
 
-    print(f"Final CTF task plan:\n{latest_plan}")
+    if latest_plan:
+        print(f"Final CTF task plan:\n{latest_plan}")
+    else:
+        print("No plan was generated due to errors.")
 
 
 if __name__ == "__main__":