from lionagi import Branch, iModel

# Pick a model
gpt4o = iModel(provider="openai", model="gpt-4o-mini")

# Create a Branch (conversation context)
hunter = Branch(
  system="you are a hilarious dragon hunter who responds in 10 words rhymes.",
  chat_model=gpt4o,
)

# Communicate asynchronously
response = await hunter.communicate("I am a dragon")
print(response)

from pydantic import BaseModel

class Joke(BaseModel):
    joke: str

res = await hunter.operate(
    instruction="Tell me a short dragon joke",
    response_format=Joke
)
print(type(res))
print(res.joke)

from lionagi.tools.types import ReaderTool

# Define model first
gpt4o = iModel(provider="openai", model="gpt-4o-mini")

branch = Branch(chat_model=gpt4o, tools=[ReaderTool])
result = await branch.ReAct(
    instruct={
      "instruction": "Summarize my PDF and compare with relevant papers.",
      "context": {"paper_file_path": "/path/to/paper.pdf"},
    },
    extension_allowed=True,     # allow multi-round expansions
    max_extensions=5,
    verbose=True,      # see step-by-step chain-of-thought
)
print(result)

from lionagi import load_mcp_tools

# Load tools from any MCP server
tools = await load_mcp_tools(".mcp.json", ["search", "memory"])

# Use with ReAct reasoning
branch = Branch(chat_model=gpt4o, tools=tools)
result = await branch.ReAct(
    instruct={"instruction": "Research recent AI developments"},
    tools=["search_exa_search"],
    max_extensions=3
)

df = branch.to_df()
print(df.tail())

from lionagi import Branch, iModel

# Define models for multi-model orchestration
gpt4o = iModel(provider="openai", model="gpt-4o-mini")
sonnet = iModel(
  provider="anthropic",
  model="claude-3-5-sonnet-20241022",
  max_tokens=1000,                    # max_tokens is required for anthropic models
)

branch = Branch(chat_model=gpt4o)
analysis = await branch.communicate("Analyze these stats", chat_model=sonnet) # Switch mid-flow

from lionagi import iModel, Branch

def create_cc_model():
  return iModel(
      provider="claude_code",
      endpoint="query_cli",
      model="sonnet",
      verbose_output=True,  # Enable detailed output for debugging
  )

# Start a coding session
orchestrator = Branch(chat_model=create_cc_model())
response = await orchestrator.communicate("Explain the architecture of protocols, operations, and branch")

# continue the session with more queries
response2 = await orchestrator.communicate("how do these parts form lionagi system")

# use structured outputs with claude code
from lionagi.fields import LIST_INSTRUCT_FIELD_MODEL, Instruct

response3 = await orchestrator.operate(
  instruct=Instruct(
    instruction="create 4 research questions for parallel discovery",
    guidance="put into `instruct_model` field as part of your structured result message",
    context="I'd like to create an orchestration system for AI agents using lionagi"
  ),
  field_models=[LIST_INSTRUCT_FIELD_MODEL],
)

len(response3.instruct_model)  # should be 4

async def handle_instruct(instruct):
  sub_branch = Branch(
    system="You are an diligent research expert.",
    chat_model=create_cc_model(),
  )
  return await sub_branch.operate(instruct=instruct)

# run in parallel across all instruct models
from lionagi.ln import alcall
responses = await alcall(response3.instruct_model, handle_instruct)

# now hand these reports back to the orchestrator
final_response = await orchestrator.communicate(
  "please synthesize these research findings into a final report",
  context=responses,
)