import json import logging import os import uuid from openai import OpenAI from infralo import Infralo # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger("infralo.complex_nesting") INFRALO_API_KEY = os.environ.get( "INFRALO_API_KEY", "your_workspace_api_key" ) INFRALO_ENDPOINT = os.environ.get("INFRALO_ENDPOINT", "http://localhost:8000") MODEL_DEPLOYMENT = os.environ.get("INFRALO_MODEL_DEPLOYMENT", "your_model_name") infralo = Infralo( api_key=INFRALO_API_KEY, endpoint=INFRALO_ENDPOINT, flush_interval=1.0, ) openai_client = OpenAI( base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1", api_key=INFRALO_API_KEY, ) # ── Tool Definitions ────────────────────────────────────────────────────────── # Schemas for LLM #1 (Root) root_tools_schema = [ { "type": "function", "function": { "name": "research_topic", "description": "Perform in-depth topic research using nested search and DB engines.", "parameters": { "type": "object", "properties": { "topic": {"type": "string", "description": "The topic to research."} }, "required": ["topic"], }, }, } ] # Schemas for LLM #2 (Nested) nested_tools_schema = [ { "type": "function", "function": { "name": "search_web", "description": "Search the web for up-to-date information on a specific query.", "parameters": { "type": "object", "properties": {"query": {"type": "string", "description": "Search query."}}, "required": ["query"], }, }, } ] # ── Real Tool Functions (Cleanly separated on top) ────────────────────────── def execute_search_web(query: str, trace, parent_span_id: str) -> str: """Search Web Tool: Executed inside nested context, child of LLM #2.""" logger.info(" Starting Search Tool...") with trace.tool( name="search_web", tool_call_id=f"call_search_{uuid.uuid4().hex[:8]}", parent_span_id=parent_span_id, # ← Explicitly link to LLM #2 ) as search_span: search_span.set_input({"query": query}) # Real search logic simulation search_result = f"Search results for query: {query}" search_span.set_output(search_result) logger.info(" Search Tool finished.") return search_result def execute_database_lookup(query: str, trace) -> dict: """Database Lookup Tool: Executed directly, automatically child of Research Tool.""" logger.info(" Research Tool: Performing DB Lookup...") db_call_id = f"call_db_{uuid.uuid4().hex[:8]}" with trace.tool(name="database_lookup", tool_call_id=db_call_id) as db_span: db_span.set_input({"query": query}) # Real database query simulation db_result = {"schema_ver": "2.4.1", "status": "active"} db_span.set_output(db_result) logger.info(" Database Tool finished.") return db_result def execute_research_topic(topic: str, trace, call_id: str) -> dict: """Research Topic Tool: Orchestrates nested LLM #2, Search Tool, and Database Tool.""" logger.info("Starting Research Tool...") with trace.tool(name="research_topic", tool_call_id=call_id) as research_span: research_span.set_input({"topic": topic}) # 1. Trigger nested LLM call (LLM #2) logger.info(" Research Tool: Triggering nested LLM Call (LLM #2)...") # trace.headers() automatically propagates the active research_topic.span_id as parent! try: raw_response_2 = openai_client.chat.completions.with_raw_response.create( model=MODEL_DEPLOYMENT, messages=[ { "role": "user", "content": f"Formulate a search query to search web about: {topic}", } ], tools=nested_tools_schema, extra_headers=trace.headers(), # x-infralo-parent-span-id = research_topic.span_id ) except Exception as exc: logger.error(" Research Tool: Nested LLM #2 failed. Details: %s", exc) research_span.set_output(f"Failed nested LLM call: {exc}") return {"status": "failed", "error": str(exc)} llm_span_id_2 = trace.update_from_response(raw_response_2) response_2 = raw_response_2.parse() logger.info(" Research Tool: LLM #2 finished. Span ID: %s", llm_span_id_2) # 2. Process tool calls returned by LLM #2 (Search Tool) nested_tool_calls = getattr(response_2.choices[0].message, "tool_calls", None) or [] search_outcomes = [] for nested_tc in nested_tool_calls: if nested_tc.function.name == "search_web": nested_args = json.loads(nested_tc.function.arguments) q = nested_args.get("query", topic) # Call search web tool (links to LLM #2) search_res = execute_search_web(q, trace, parent_span_id=llm_span_id_2) search_outcomes.append(search_res) # 3. Perform database lookup directly (links to Research Tool) db_res = execute_database_lookup("SELECT schema_ver FROM obs_config", trace) outcome = { "status": "completed", "nested_summary": "Deep research complete.", "search_count": len(search_outcomes), "db_status": db_res.get("status"), } research_span.set_output(outcome) logger.info("Research Tool finished.") return outcome # ── Main Orchestration ──────────────────────────────────────────────────────── def run_complex_nesting(): print("=== Running Complex Multi-Nesting Trace Example ===") with infralo.start_trace(session_id="session_complex_nesting") as trace: print(f"Started root trace: {trace.trace_id}") # 1. ROOT LLM Call (LLM #1) prompt = "Please perform deep research on the topic 'observability engines'." logger.info("Executing Root LLM Call (LLM #1)...") try: raw_response_1 = openai_client.chat.completions.with_raw_response.create( model=MODEL_DEPLOYMENT, messages=[{"role": "user", "content": prompt}], tools=root_tools_schema, extra_headers=trace.headers(), ) except Exception as exc: logger.error("Failed to connect to Gateway (LLM #1). Details: %s", exc) return llm_span_id_1 = trace.update_from_response(raw_response_1) response_1 = raw_response_1.parse() logger.info("LLM #1 finished. Span ID: %s", llm_span_id_1) # 2. Iterate and execute tool calls tool_calls = getattr(response_1.choices[0].message, "tool_calls", None) or [] research_outcome = {} for tool_call in tool_calls: if tool_call.function.name == "research_topic": args = json.loads(tool_call.function.arguments) topic = args["topic"] call_id = tool_call.id # Execute the real top-level tool function research_outcome = execute_research_topic(topic, trace, call_id) # 3. Finalize root LLM call with tool outputs # trace.headers() only carries x-infralo-trace-id here (no active tool span), # so this final LLM call is a root-level sibling — not a child of research_topic. logger.info("Sending final outputs to LLM #1...") try: raw_response_final = openai_client.chat.completions.with_raw_response.create( model=MODEL_DEPLOYMENT, messages=[ {"role": "user", "content": prompt}, response_1.choices[0].message, { "role": "tool", "tool_call_id": tool_calls[0].id, "name": "research_topic", "content": json.dumps(research_outcome), }, ], extra_headers=trace.headers(), ) final_resp = raw_response_final.parse() print(f"Final LLM Response: {final_resp.choices[0].message.content}") except Exception as exc: logger.error("Final LLM call failed. Details: %s", exc) if __name__ == "__main__": run_complex_nesting()