Hello, I am trying to create an agent that will call a tool and fetch me some details from an API. If I use a model from OpenAI the tool gets called and everything works perfectly but when I try to run the same using models from Ollama, the agents keep giving response form it’s memory.
import os
import requests
from datetime import datetime, timedelta, timezone
from typing import List, Callable, Dict, Any, Optional
from typing import Dict, Iterator, Optional
from agno.agent import Agent
from agno.models.ollama import Ollama
from agno.models.openai import OpenAIChat
from pydantic import BaseModel, Field
from constants import MODEL_ID
def get_model():
"""Function to get the model."""
# model = Ollama(id=MODEL_ID)
model = OpenAIChat(id="gpt-4o")
print(MODEL_ID)
return model
def model_wrapper(func):
"""Decorator to wrap the agent with the model."""
def wrapped(*args, **kwargs):
agent = func(*args, **kwargs)
agent.model = get_model()
return agent
return wrapped
class VulnerabilitiesDesign(BaseModel):
cveNumber : str = Field(..., description="CVE identifier for the vulnerability.")
cveTitle : str = Field(..., description="Title of the CVE.")
description: str = Field(..., description="Description of the vulnerability.")
severity: str = Field(..., description="Severity of the vulnerability.")
impact: str = Field(..., description="Impact of the vulnerability.")
AffectedSystem: str = Field(..., description="Systems or applications affected by the vulnerability.")
Url: str = Field(..., description="URL to the CVE advisory or details.")
class VulnerabilityList(BaseModel):
Vulnerabilities: list[VulnerabilitiesDesign]
def is_within_last_n_days(release_date_str: str, n:int) -> bool:
release_date = datetime.strptime(release_date_str, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
now = datetime.now(timezone.utc)
return now - timedelta(days=n) <= release_date <= now
def fetch_vulnerability():
"""Use this function to fetch vulnerabilityb related to Windows OS and applications.
Args:
num (int): Number of days to look back for vulnerabilities. Default is 3.
Returns:
List of vulnerabilities found in the last `num` days."""
print("Entered fetch_cve")
BASE_URL = "https://api.msrc.microsoft.com/sug/v2.0/en-us/vulnerability"
num = 3
HEADERS = {
"Accept": "application/json",
}
url = BASE_URL
response = requests.get(url, headers=HEADERS)
print("Status Code:", response.status_code)
print("URL:", response.url)
response.raise_for_status()
updates = response.json()
result = []
for update in updates.get("value", []):
count = 0
if is_within_last_n_days(update['releaseDate'], num):
vulnerability = check_severity(update, "Critical")
if vulnerability:
count += 1
result.append(vulnerability)
if count >= 10:
break
if not result: # If no critical vulnerabilities found
for update in updates.get("value", []):
if is_within_last_n_days(update['releaseDate'], num):
vulnerability = check_severity(update, "Important")
if vulnerability:
count += 1
result.append(vulnerability)
if count >= 10:
break
print(len(result))
return result
def check_severity(update, severity) -> Dict[str, Any]:
if update['severity'] == severity:
result = {
'cveNumber': update.get('cveNumber', None),
'cveTitle': update.get('cveTitle', None),
'description': update.get('description', None),
'severity': update.get('severity', None),
'impact': update.get('impact', None),
'AffectedSystem': update.get('tag', None),
'Url': update.get('mitreUrl', None),
}
return result
# Defining the different agents
@model_wrapper
def create_windows_cve_agent():
windows_agent_instance = Agent(
model = None,
# context={"windows vulnerability": fetch_cve},
# enable_thinking = True,
# role="Fetch vulnerabilities in the Windows OS and applications using tools.",
description=dedent("""\
You are an agent that extracts structured security vulnerability details.\
Your output should only include specific vulnerabilities and description \
(no general summaries). The tool helps you fetch the latest vulnerabilities in Windows OS and applications.
"""),
tools=[fetch_vulnerability],
# tools = [Windows_Vulnerability_search()],
# tool_choice={"type": "function", "function": {"name": "fetch_vulnerabilities"}},
tool_choice = 'auto',
debug_mode=True,
show_tool_calls=True,
response_model=VulnerabilityList,)
# structured_outputs=True,)
# use_json_mode=True)
return windows_agent_instance
windows_cve_agent = create_windows_cve_agent()
if __name__ == "__main__":
windows_cve_agent.print_response("Top 5 Windows vulnerabilities from the past 3 days")
I have tried running using qwen3:14b, mistral-small3.1:24b, cogito:32b