kitai

Imports

from cachy import enable_cachy
enable_cachy()

m = "anthropic/claude-haiku-4-5"

Model response helpers

assistant_response


def assistant_response(
    content:NoneType=None, tool_calls:NoneType=None, finish_reason:str='stop'
):

Call self as a function.

tool_response


def tool_response(
    name:str, arguments:dict
):

Call self as a function.

mk_cb_msg


def mk_cb_msg(
    msg:str | functools.partial | litellm.types.utils.ModelResponse
):

Call self as a function.

drain_result


def drain_result(
    gen
):

Call self as a function.

Tool helpers

mk_schema


def mk_schema(
    fn, strip:tuple=('tool_ctx',)
):

Call self as a function.

call_llm_tool


def call_llm_tool(
    tc_func:dict, tool_map:dict, state:dict
):

Call self as a function.

Streaming completions

accumulate_toolcalls


def accumulate_toolcalls(
    tool_buf:dict, delta
):

Call self as a function.

stream_completion


def stream_completion(
    model, messages, tools:NoneType=None, agent:str='', kwargs:VAR_KEYWORD
):

Call self as a function.

Message preprocessing

msg_fragment


def msg_fragment(
    d:dict, results:dict
):

Call self as a function.

mk_narrative_cast


def mk_narrative_cast(
    msg_list:list, active_agent:str
):

Call self as a function.

filter_out_history


def filter_out_history(
    msgs:list, active_agent:str
):

Call self as a function.

Agent


def Agent(
    model:str, # model str in litellm format
    name:str='', # name of the agent
    desc:str='', # what the agent does, used by parent agents to decide when to transfer
    sp:NoneType=None, # system prompt
    tools:NoneType=None, # tools provided to the agent
    subagents:NoneType=None, # subagents this agent can transfer to
    narrative_cast:bool=True, # transform prior agent messages into narrative form
    include_hist:bool=True, # if False, remove older history not belonging to this agent
    before_cb:NoneType=None, # callback func(state, msgs); non-None return skips LLM
    after_cb:NoneType=None, # callback func(state, response); can replace LLM response
    output_key:NoneType=None, # save model response content to state[output_key]
    llm_kwargs:NoneType=None
):

Initialize self. See help(type(self)) for accurate signature.

Agent.preproc_msg


def preproc_msg(
    msgs:list | str, state:NoneType=None
):

Call self as a function.

Agent.run_callback


def run_callback(
    fn, msgs:list, state:dict
):

Call self as a function.

Agent.maybe_save_to_output_key


def maybe_save_to_output_key(
    state, response
):

Call self as a function.

Agent.call


def __call__(
    msgs:list | str, state:NoneType=None, kwargs:VAR_KEYWORD
):

Call self as a function.

Agent.stream


def stream(
    msgs:list | str, state:NoneType=None, kwargs:VAR_KEYWORD
):

Call self as a function.

Runner

with_model_nm


def with_model_nm(
    llm_output, nm:str
):

Call self as a function.

mk_transfer_doc


def mk_transfer_doc(
    agent, parent
):

Call self as a function.

Runner


def Runner(
    agent:Agent
):

Initialize self. See help(type(self)) for accurate signature.

Sequential


def Sequential(
    agent_list:list
):

Initialize self. See help(type(self)) for accurate signature.

Examples

def mult(a:int, b:int):
    "Multiplies two numbers together"
    return a * b

def save(result:str, tool_ctx):
    "Saves some result passed into storage"
    tool_ctx["final_result"] = result

saver = Agent(m, name="saver", desc="Saves results into storage.", sp="You receive results from other agents and save them with your tool.", tools=[save])
math_agent = Agent(m, name="math_agent", desc="Is good at doing maths.", sp="Always do maths using tools and always pass to saver to save results.", tools=[mult], subagents=[saver])
root = Agent(m, name="root", desc="General agent", sp="A general coordinator that delegates to sub agents as needed", subagents=[math_agent])

r = Runner(root)
r("What's 359231 x 235981")

Tool calls: - 🔧 transfer_to_agent {"agent_name": "math_agent"}

Tool calls: - 🔧 mult {"a": 359231, "b": 235981}

Tool calls: - 🔧 transfer_to_agent {"agent_name": "saver"}

Tool calls: - 🔧 save {"result": "359231 × 235981 = 84,771,690,611"}

[{'content': None,
  'role': 'assistant',
  'tool_calls': [{'function': {'arguments': '{"agent_name": "math_agent"}',
     'name': 'transfer_to_agent'},
    'id': 'toolu_0128PhvCTy9S8CjWFvHWGEx6',
    'type': 'function'}],
  'function_call': None,
  'provider_specific_fields': None,
  'agent': 'root'},
 {'role': 'tool',
  'tool_call_id': 'toolu_0128PhvCTy9S8CjWFvHWGEx6',
  'content': 'null',
  'agent': 'root'},
 {'content': "I'll calculate 359231 × 235981 for you using the multiplication tool, and then transfer the result to the saver agent.",
  'role': 'assistant',
  'tool_calls': [{'function': {'arguments': '{"a": 359231, "b": 235981}',
     'name': 'mult'},
    'id': 'toolu_01HwBAkHnpdh1ASVsqDHf9qo',
    'type': 'function'}],
  'function_call': None,
  'provider_specific_fields': None,
  'agent': 'math_agent'},
 {'role': 'tool',
  'tool_call_id': 'toolu_01HwBAkHnpdh1ASVsqDHf9qo',
  'content': '84771690611',
  'agent': 'math_agent'},
 {'content': 'Great! The result is **84,771,690,611**. Now let me transfer this to the saver agent to save the result:',
  'role': 'assistant',
  'tool_calls': [{'function': {'arguments': '{"agent_name": "saver"}',
     'name': 'transfer_to_agent'},
    'id': 'toolu_01M4aC8gLJcjJZ292pQfM9W9',
    'type': 'function'}],
  'function_call': None,
  'provider_specific_fields': None,
  'agent': 'math_agent'},
 {'role': 'tool',
  'tool_call_id': 'toolu_01M4aC8gLJcjJZ292pQfM9W9',
  'content': 'null',
  'agent': 'math_agent'},
 {'content': "I'll save that result for you.",
  'role': 'assistant',
  'tool_calls': [{'function': {'arguments': '{"result": "359231 × 235981 = 84,771,690,611"}',
     'name': 'save'},
    'id': 'toolu_01XA5i1ySEcQFhQXEUw4A4nf',
    'type': 'function'}],
  'function_call': None,
  'provider_specific_fields': None,
  'agent': 'saver'},
 {'role': 'tool',
  'tool_call_id': 'toolu_01XA5i1ySEcQFhQXEUw4A4nf',
  'content': 'null',
  'agent': 'saver'},
 {'content': 'The calculation has been completed and saved! **359231 × 235981 = 84,771,690,611**',
  'role': 'assistant',
  'tool_calls': None,
  'function_call': None,
  'provider_specific_fields': None,
  'agent': 'saver'}]

a = Agent(m, name="test", sp="simple test", tools=[mult])
r = Runner(a)
for event in r.stream("First give a small intro, then multiply 598235 * 2983535 and give me the result"):
    print(event)

{'role': 'assistant', 'content': 'I have', 'agent': 'test'}
{'role': 'assistant', 'content': ' access to a multiplication function that can help with your calculation. Let me use it to multiply those two large numbers for', 'agent': 'test'}
{'role': 'assistant', 'content': ' you.', 'agent': 'test'}
{'role': 'tool', 'tool_call_id': 'toolu_01K3AQhoKgYZVwhAcpizniFD', 'content': '1784855060725', 'agent': 'test', 'function': {'arguments': '{"a": 598235, "b": 2983535}', 'name': 'mult'}}
{'role': 'assistant', 'content': "Here's your", 'agent': 'test'}
{'role': 'assistant', 'content': " result:\n\n**598235 × 2,983,535 = 1,784,855,060,725**\n\nThat's over 1", 'agent': 'test'}
{'role': 'assistant', 'content': ".7 trillion! These are fairly large numbers, so it's easy", 'agent': 'test'}
{'role': 'assistant', 'content': ' to see why having a tool to handle the multiplication accurately is helpful.', 'agent': 'test'}

Local models

You can inference on a local model (e.g., vLLM hoste) by creating the agent like:

a = Agent("openai/model_name", llm_kwargs={"api_base": "http://localhost:8000/v1", "api_key": "none"})

Replacing model_name with the name provided when launching vLLM.

Imports

Model response helpers

assistant_response

tool_response

mk_cb_msg

drain_result

Tool helpers

mk_schema

call_llm_tool

Streaming completions

accumulate_toolcalls

stream_completion

Message preprocessing

msg_fragment

mk_narrative_cast

filter_out_history

Agent

Agent

Agent.preproc_msg

Agent.run_callback

Agent.maybe_save_to_output_key

Agent.__call__

Agent.stream

Runner

with_model_nm

mk_transfer_doc

Runner

Sequential

Sequential

Examples

Local models

Agent.call