Highlighting the integration of Ollama, OpenAI Swarm, and Pydantic to create structured AI agents. Features keywords, data streams, and a developer interface symbolizing local-first AI frameworks and structured programming practices.

swarm_ollama_pydantic.py

from openai import OpenAI
from rich import print

client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

# model="qwq"
model="qwen2.5-coder:32b"

# TESTING INFERENCE

# response = client.chat.completions.create(
#     model=model,
#     n=1,
#     max_tokens=100,
#     messages=[
#         {"role": "system", "content": "You are a helpful assistant."},
#         {"role": "user", "content": "Explain how AI works"}
#     ]
# )

# print(response.choices[0].message)

# Initialize Swarm
from swarm import Swarm, Agent
from swarm.types import Response
from openai import OpenAI
from rich import print
import os, re, json
from pydantic import BaseModel
from typing import List, Optional, Dict, Any

# Define our Pydantic class to go with the structured output model
class PersonInfo(BaseModel):
    name: str
    age: int
    skills: List[str]
    bio: Optional[str] = None

def clean_json_string(data_str: str) -> str:
    """Clean JSON string by removing markdown code blocks and other formatting"""
    # Remove markdown code blocks if present
    if "```" in data_str:
        print("Removing markdown code blocks...")
        # Extract content between code blocks
        match = re.search(r'```(?:json)?\n(.*?)\n```', data_str, re.DOTALL)
        if match:
            data_str = match.group(1)
    return data_str.strip()


def process_extracted_data(name: str, age: int, skills: List[str], bio: Optional[str] = None) -> Optional[PersonInfo]:
    """Helper function to process extracted data"""
    print("\n=== process_extracted_data called ===")
    print(f"Received data:")
    print(f"name: {name}")
    print(f"age: {age}")
    print(f"skills: {skills}")
    print(f"bio: {bio}")
    
    try:
        result = PersonInfo(
            name=name,
            age=age,
            skills=skills,
            bio=bio
        )
        print(f"Successfully created PersonInfo: {result}")
        return result
    except Exception as e:
        print(f"Error processing data: {e}")
        return None
    finally:
        print("=== process_extracted_data finished ===\n")

def create_person_info_agent() -> Agent:
    return Agent(
        name="PersonInfoAgent",
        instructions="""You are a precise information extraction agent that converts unstructured text about people into a specific JSON format.

IMPORTANT: When calling process_extracted_data, you MUST format the data exactly as follows:
{
    "name": "string",
    "age": number,
    "skills": ["skill1", "skill2"],  # MUST be a JSON array/list of strings
    "bio": "string"
}

The skills parameter MUST ALWAYS be a JSON array/list of strings, NOT a comma-separated string.

CORRECT format for skills:
  "skills": ["AI Agents", "RAG workflows"]

INCORRECT format for skills:
  "skills": "AI Agents, RAG workflows"

Example input: "John Smith is a 35-year-old software developer skilled in Python and Cloud Architecture."
You should call process_extracted_data with:
{
    "name": "John Smith",
    "age": 35,
    "skills": ["Python", "Cloud Architecture"],
    "bio": "Software developer"
}""",
        functions=[process_extracted_data]
    )

agent = create_person_info_agent()

test_case = {
    "text": "Pat Lesieur is a 65-year-old software developer skilled in AI Agents and RAG worflows."
}

# Initialize Swarm
from swarm import Agent, Swarm
swarm_client = Swarm(client=client)

# Run the agent with the test input
response = swarm_client.run(
    agent=agent,
    model_override=model,
    messages=[{
        "role": "user", 
        "content": test_case['text']
    }],
    context_variables={},
    max_turns=1,
    execute_tools=True  # Make sure this is True to actually execute the function
)

print("\n=== Complete Response Details ===")
for msg in response.messages:
    print(f"\nMessage type: {msg['role']}")
    print(f"Content: {msg.get('content')}")
    if msg.get('function_call'):
        print(f"Function call: {json.dumps(msg['function_call'], indent=2)}")
    if msg.get('tool_calls'):
        print(f"Tool calls: {json.dumps(msg['tool_calls'], indent=2)}")