From 3134c572c76479224189fea51231fa6c0aada772 Mon Sep 17 00:00:00 2001 From: etserend Date: Mon, 13 Jul 2026 18:21:38 -0500 Subject: [PATCH 1/2] Revert "feat(tooling): bump ruff 0.9.7 -> 0.15.20 and openapi-python-client 0.26.1 -> 0.29.0 (HYBIM-777) (#66)" This reverts commit e121f413e884223545b5acf3d6affde2e6a9a0ce. --- .pre-commit-config.yaml | 2 +- .../research_crew/src/research_crew/crew.py | 15 +- .../research_crew/src/research_crew/main.py | 6 +- .../src/research_crew/tools/custom_tool.py | 7 +- examples/agent/google-adk/my_agent/agent.py | 9 +- examples/agent/langchain-agent/main.py | 3 +- examples/agent/langchain-middleware/main.py | 20 +-- .../after/create_sample_logs.py | 1 - .../after/scripts/setup_pinecone.py | 4 +- .../agent/langgraph-fsi-agent/after/test.py | 7 +- .../before/scripts/setup_pinecone.py | 4 +- .../agent/langgraph-open-telemetry/main.py | 10 +- examples/agent/langgraph-otel/main.py | 30 ++-- examples/agent/langgraph-telecom-agent/app.py | 4 +- .../scripts/setup_pinecone.py | 4 +- .../src/tools/billing_tool.py | 24 +-- .../src/tools/technical_support_tool.py | 2 +- examples/agent/langgraph-traceloop/agent.py | 35 ++-- examples/agent/langgraph-traceloop/main.py | 9 +- .../agent/microsoft-agent-framework/agent.py | 9 +- examples/agent/minimal-agent-example/app.py | 93 ++++++---- .../minimal-agent-example/minimal_agent.py | 37 ++-- .../pydantic-ai-support-agent/src/main.py | 54 ++++-- .../agent/startup-simulator-3000/__init__.py | 7 +- .../agent/startup-simulator-3000/agent.py | 86 +++++----- .../agent_framework/__init__.py | 8 +- .../agent_framework/agent.py | 89 ++++++---- .../agent_framework/config.py | 27 ++- .../agent_framework/exceptions.py | 14 +- .../agent_framework/factory.py | 22 ++- .../agent_framework/llm/__init__.py | 4 +- .../agent_framework/llm/base.py | 9 +- .../agent_framework/llm/models.py | 23 ++- .../agent_framework/llm/openai_provider.py | 47 +++-- .../agent_framework/llm/tool_models.py | 36 ++-- .../agent_framework/models.py | 160 +++++++++--------- .../agent_framework/prompts/templates.py | 9 +- .../agent_framework/state.py | 12 +- .../agent_framework/tools/base.py | 8 +- .../agent_framework/tools/registry.py | 17 +- .../agent_framework/utils/formatting.py | 38 +++-- .../agent_framework/utils/hooks.py | 31 ++-- .../agent_framework/utils/logging.py | 34 ++-- .../agent_framework/utils/tool_hooks.py | 34 +++- .../agent_framework/utils/tool_registry.py | 21 ++- .../agent_framework/utils/validation.py | 5 +- examples/agent/startup-simulator-3000/app.py | 23 +-- examples/agent/startup-simulator-3000/demo.py | 15 +- .../startup-simulator-3000/run_startup_sim.py | 10 +- .../startup-simulator-3000/tools/__init__.py | 6 +- .../tools/hackernews_tool.py | 35 ++-- .../tools/keyword_extraction.py | 24 ++- .../tools/news_api_tool.py | 115 ++++++++----- .../tools/serious_startup_simulator.py | 60 ++++--- .../tools/startup_simulator.py | 48 +++--- .../tools/text_analysis.py | 7 +- examples/agent/strands-agents/agent.py | 9 +- .../agent/weather-vibes-agent/__init__.py | 4 +- examples/agent/weather-vibes-agent/agent.py | 62 +++++-- .../weather-vibes-agent/agent/descriptor.py | 12 +- .../agent/weather_vibes_agent.py | 66 +++++--- .../agent/weather-vibes-agent/descriptor.py | 12 +- .../weather-vibes-agent/tools/__init__.py | 4 +- .../tools/recommendation_tool.py | 9 +- .../weather-vibes-agent/tools/weather_tool.py | 14 +- .../weather-vibes-agent/tools/youtube_tool.py | 17 +- examples/chatbot/basic-examples/app.py | 13 +- examples/chatbot/basic-examples/context.py | 6 +- .../chatbot/basic-examples/hallucination.py | 19 ++- examples/chatbot/basic-examples/test.py | 15 +- examples/chatbot/elevenlabs-chatbot/main.py | 2 +- .../elevenlabs-chatbot/splunk_ao_handler.py | 20 ++- .../sample-project-chatbot/anthropic/app.py | 10 +- .../sample-project-chatbot/anthropic/test.py | 26 ++- .../azure-inference/app.py | 7 +- .../azure-inference/test.py | 26 ++- .../openai-ollama/app.py | 4 +- .../openai-ollama/create_sample_logs.py | 10 +- .../openai-ollama/test.py | 26 ++- examples/dataset-experiments/app-simple.py | 7 +- examples/dataset-experiments/app.py | 36 ++-- .../experiment_compare_two_models.py | 55 +++--- examples/experiments/rag-and-tools/app.py | 37 +++- .../experiments/rag-and-tools/experiment.py | 13 +- .../upload_existing_results.py | 44 +++-- .../python-service/app.py | 38 +++-- .../python-service/seed_chroma.py | 17 +- .../distributed-tracing/main_run.py | 24 +-- .../distributed-tracing/retrieval_service.py | 7 +- .../galileologger/basic-example.py | 6 +- .../galileologger/metadata-example.py | 7 +- .../galileologger/retriever-example.py | 14 +- examples/logging-samples/log-mcp-calls/app.py | 15 +- .../log-mcp-calls/mcp_client.py | 24 ++- .../logging-samples/openai-responses/main.py | 18 +- examples/rag/cli-rag-demo/app.py | 29 ++-- .../challenges/chunk-utilization.py | 108 ++++++++---- .../challenges/chunk-utilizations-basic.py | 17 +- .../challenges/chunk-utilizations-enhanced.py | 25 +-- .../cli-rag-demo/challenges/document_store.py | 67 ++++---- .../challenges/document_store_basic.py | 20 +-- .../challenges/ensure-completeness-basic.py | 19 ++- .../ensure-completeness-enhanced.py | 19 ++- .../cli-rag-demo/challenges/out-of-context.py | 41 +++-- examples/rag/cli-rag-demo/test.py | 7 +- .../rag/elastic-chatbot-rag-app/api/app.py | 2 +- .../rag/elastic-chatbot-rag-app/api/chat.py | 38 ++++- .../api/elasticsearch_client.py | 3 +- .../api/llm_integrations.py | 5 +- .../rag/elastic-chatbot-rag-app/api/utils.py | 4 +- .../data/index_data.py | 32 ++-- poetry.lock | 128 ++++++-------- pyproject.toml | 21 +-- scripts/create_docs.py | 7 +- splunk-ao-a2a/examples/two_agent_demo.py | 2 +- splunk-ao-a2a/src/splunk_ao_a2a/_context.py | 2 +- splunk-ao-adk/src/splunk_ao_adk/callback.py | 1 + splunk-ao-adk/src/splunk_ao_adk/observer.py | 1 + splunk-ao-adk/src/splunk_ao_adk/plugin.py | 1 + .../src/splunk_ao_adk/span_manager.py | 1 + .../src/splunk_ao_adk/trace_builder.py | 1 - splunk-ao-adk/tests/conftest.py | 3 +- splunk-ao-adk/tests/test_retriever.py | 1 - splunk-ao-adk/tests/test_trace_builder.py | 2 +- src/splunk_ao/__init__.py | 4 +- src/splunk_ao/constants/routes.py | 4 +- src/splunk_ao/experiment.py | 2 +- src/splunk_ao/experiments.py | 8 +- .../handlers/agent_control/bridge.py | 6 +- src/splunk_ao/handlers/base_handler.py | 4 +- .../handlers/openai_agents/handler.py | 8 +- src/splunk_ao/logger/control.py | 12 +- src/splunk_ao/schema/handlers.py | 4 +- src/splunk_ao/schema/message.py | 3 - src/splunk_ao/schema/metrics.py | 4 +- src/splunk_ao/schema/trace.py | 8 +- src/splunk_ao/search.py | 4 +- src/splunk_ao/utils/__init__.py | 4 +- src/splunk_ao/utils/serialization.py | 2 +- tests/test_agent_control_bridge.py | 2 +- tests/test_base_handler.py | 4 +- tests/test_config.py | 8 +- tests/test_configuration.py | 4 +- tests/test_crewai_handler.py | 68 ++++---- tests/test_decorator.py | 2 +- tests/test_decorator_distributed.py | 2 +- tests/test_experiment.py | 16 +- tests/test_experiment_tags.py | 4 +- tests/test_experiments.py | 27 +-- tests/test_langchain.py | 8 +- tests/test_langchain_async.py | 12 +- tests/test_langchain_middleware.py | 8 +- tests/test_log_streams_metrics.py | 4 +- tests/test_logger_batch.py | 2 +- tests/test_logger_distributed.py | 2 + tests/test_logger_timestamps.py | 8 +- tests/test_openai.py | 2 +- tests/utils/test_datasets.py | 4 +- tests/utils/test_serialization.py | 4 +- 159 files changed, 1784 insertions(+), 1319 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index fb3732ae..46c086e6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -11,7 +11,7 @@ repos: - id: end-of-file-fixer - id: trailing-whitespace - repo: https://github.com/charliermarsh/ruff-pre-commit - rev: "v0.15.20" + rev: "v0.9.7" hooks: - id: ruff args: [--fix] diff --git a/examples/agent/crewAI/research_crew/src/research_crew/crew.py b/examples/agent/crewAI/research_crew/src/research_crew/crew.py index c8d5a859..b878e27d 100644 --- a/examples/agent/crewAI/research_crew/src/research_crew/crew.py +++ b/examples/agent/crewAI/research_crew/src/research_crew/crew.py @@ -1,17 +1,17 @@ # src/research_crew/crew.py - from crewai import Agent, Crew, Process, Task -from crewai.agents.agent_builder.base_agent import BaseAgent from crewai.project import CrewBase, agent, crew, task from crewai_tools import SerperDevTool +from crewai.agents.agent_builder.base_agent import BaseAgent +from typing import List @CrewBase class ResearchCrew: """Research crew for comprehensive topic analysis and reporting""" - agents: list[BaseAgent] - tasks: list[Task] + agents: List[BaseAgent] + tasks: List[Task] @agent def researcher(self) -> Agent: @@ -32,4 +32,9 @@ def analysis_task(self) -> Task: @crew def crew(self) -> Crew: """Creates the research crew""" - return Crew(agents=self.agents, tasks=self.tasks, process=Process.sequential, verbose=True) + return Crew( + agents=self.agents, + tasks=self.tasks, + process=Process.sequential, + verbose=True, + ) diff --git a/examples/agent/crewAI/research_crew/src/research_crew/main.py b/examples/agent/crewAI/research_crew/src/research_crew/main.py index aef88355..ece1c5d8 100644 --- a/examples/agent/crewAI/research_crew/src/research_crew/main.py +++ b/examples/agent/crewAI/research_crew/src/research_crew/main.py @@ -1,10 +1,8 @@ #!/usr/bin/env python # src/research_crew/main.py import os - -from dotenv import load_dotenv - from research_crew.crew import ResearchCrew +from dotenv import load_dotenv from splunk_ao.handlers.crewai.handler import CrewAIEventListener load_dotenv() @@ -13,7 +11,7 @@ os.makedirs("output", exist_ok=True) -def run() -> None: +def run(): # Create the event listener for Splunk AO CrewAI integration CrewAIEventListener() diff --git a/examples/agent/crewAI/research_crew/src/research_crew/tools/custom_tool.py b/examples/agent/crewAI/research_crew/src/research_crew/tools/custom_tool.py index ec9fd9cd..08739e65 100644 --- a/examples/agent/crewAI/research_crew/src/research_crew/tools/custom_tool.py +++ b/examples/agent/crewAI/research_crew/src/research_crew/tools/custom_tool.py @@ -1,4 +1,5 @@ from crewai.tools import BaseTool +from typing import Type from pydantic import BaseModel, Field @@ -10,10 +11,8 @@ class MyCustomToolInput(BaseModel): class MyCustomTool(BaseTool): name: str = "Name of my tool" - description: str = ( - "Clear description for what this tool is useful for, your agent will need this information to use it." - ) - args_schema: type[BaseModel] = MyCustomToolInput + description: str = "Clear description for what this tool is useful for, your agent will need this information to use it." + args_schema: Type[BaseModel] = MyCustomToolInput def _run(self, argument: str) -> str: # Implementation goes here diff --git a/examples/agent/google-adk/my_agent/agent.py b/examples/agent/google-adk/my_agent/agent.py index 6342e9da..438418cd 100644 --- a/examples/agent/google-adk/my_agent/agent.py +++ b/examples/agent/google-adk/my_agent/agent.py @@ -1,11 +1,10 @@ """Agent logic for the Google ADK example.""" from dotenv import load_dotenv -from google.adk.agents.llm_agent import Agent -from openinference.instrumentation.google_adk import GoogleADKInstrumentor from opentelemetry.sdk import trace as trace_sdk - from splunk_ao import otel +from openinference.instrumentation.google_adk import GoogleADKInstrumentor +from google.adk.agents.llm_agent import Agent load_dotenv() @@ -32,8 +31,6 @@ def get_current_time(city: str) -> dict: model="gemini-3-flash-preview", name="root_agent", description="Tells the current time in a specified city.", - instruction=( - "You are a helpful assistant that tells the current time in cities.Use the 'get_current_time' tool for this purpose." - ), + instruction=("You are a helpful assistant that tells the current time in cities." "Use the 'get_current_time' tool for this purpose."), tools=[get_current_time], ) diff --git a/examples/agent/langchain-agent/main.py b/examples/agent/langchain-agent/main.py index b69c5a65..5da1302a 100644 --- a/examples/agent/langchain-agent/main.py +++ b/examples/agent/langchain-agent/main.py @@ -1,9 +1,8 @@ from dotenv import load_dotenv from langchain.agents import initialize_agent from langchain.agents.agent_types import AgentType -from langchain.tools import tool from langchain_openai import ChatOpenAI - +from langchain.tools import tool from splunk_ao import splunk_ao_context from splunk_ao.handlers.langchain import SplunkAOCallback diff --git a/examples/agent/langchain-middleware/main.py b/examples/agent/langchain-middleware/main.py index 3fd8cef3..a50a29a9 100644 --- a/examples/agent/langchain-middleware/main.py +++ b/examples/agent/langchain-middleware/main.py @@ -2,14 +2,13 @@ import os from dotenv import load_dotenv +from splunk_ao import splunk_ao_context +from splunk_ao.handlers.langchain.middleware import SplunkAOMiddleware from langchain.agents.factory import create_agent from langchain.tools import tool from langchain_core.messages import HumanMessage from langchain_openai import ChatOpenAI -from splunk_ao import splunk_ao_context -from splunk_ao.handlers.langchain.middleware import SplunkAOMiddleware - # Load environment variables (e.g., API keys) load_dotenv() @@ -43,10 +42,7 @@ def get_stock_price(symbol: str) -> str: def main() -> None: # Use the Splunk AO context manager to specify project and log stream # All traces created within this context will be associated with this project - with splunk_ao_context( - project=os.getenv("SPLUNK_AO_PROJECT", "langchain-middleware"), - log_stream=os.getenv("SPLUNK_AO_LOG_STREAM", "agent_execution"), - ): + with splunk_ao_context(project=os.getenv("SPLUNK_AO_PROJECT", "langchain-middleware"), log_stream=os.getenv("SPLUNK_AO_LOG_STREAM", "agent_execution")): # Create an agent with SplunkAOMiddleware for automatic logging # SplunkAOMiddleware automatically captures: # - Agent lifecycle events (start/completion) @@ -65,15 +61,7 @@ def main() -> None: # 3. Understand it needs Apple's stock price # 4. Call the get_stock_price tool # 5. Synthesize the results into a coherent response - result = agent.invoke( - { - "messages": [ - HumanMessage( - content="What's the weather like in San Francisco and what's the current stock price of Apple?" - ) - ] - } - ) + result = agent.invoke({"messages": [HumanMessage(content="What's the weather like in San Francisco and what's the current stock price of Apple?")]}) print(f"\nAgent Response:\n{result}") diff --git a/examples/agent/langgraph-fsi-agent/after/create_sample_logs.py b/examples/agent/langgraph-fsi-agent/after/create_sample_logs.py index 4de35455..61640cba 100644 --- a/examples/agent/langgraph-fsi-agent/after/create_sample_logs.py +++ b/examples/agent/langgraph-fsi-agent/after/create_sample_logs.py @@ -1,5 +1,4 @@ # A script to generate log streams -# ruff: noqa: E402 -- load_dotenv() must precede all other imports from dotenv import load_dotenv diff --git a/examples/agent/langgraph-fsi-agent/after/scripts/setup_pinecone.py b/examples/agent/langgraph-fsi-agent/after/scripts/setup_pinecone.py index b53d7bdc..4f3b81fe 100644 --- a/examples/agent/langgraph-fsi-agent/after/scripts/setup_pinecone.py +++ b/examples/agent/langgraph-fsi-agent/after/scripts/setup_pinecone.py @@ -101,9 +101,7 @@ def upload_to_pinecone(chunked_docs, index_name: str, force_upload: bool = False # Create vector store and upload print(f"Uploading {len(chunked_docs)} chunks to Pinecone...") - vector_store = PineconeVectorStore.from_documents( - documents=chunked_docs, embedding=EMBEDDINGS, index_name=index_name - ) + vector_store = PineconeVectorStore.from_documents(documents=chunked_docs, embedding=EMBEDDINGS, index_name=index_name) print(f"Successfully uploaded {len(chunked_docs)} document chunks to Pinecone") return vector_store diff --git a/examples/agent/langgraph-fsi-agent/after/test.py b/examples/agent/langgraph-fsi-agent/after/test.py index d8bb2c6b..3fcf41b0 100644 --- a/examples/agent/langgraph-fsi-agent/after/test.py +++ b/examples/agent/langgraph-fsi-agent/after/test.py @@ -108,12 +108,7 @@ def test_run_experiment_with_dataset(): experiment_name="langgraph-fsi-experiment", dataset_name=DATASET_NAME, function=send_message_to_supervisor_agent, - metrics=[ - SplunkAOMetrics.action_advancement, - SplunkAOMetrics.action_completion, - SplunkAOMetrics.tool_error_rate, - SplunkAOMetrics.tool_selection_quality, - ], + metrics=[SplunkAOMetrics.action_advancement, SplunkAOMetrics.action_completion, SplunkAOMetrics.tool_error_rate, SplunkAOMetrics.tool_selection_quality], project=os.getenv("SPLUNK_AO_PROJECT"), ) diff --git a/examples/agent/langgraph-fsi-agent/before/scripts/setup_pinecone.py b/examples/agent/langgraph-fsi-agent/before/scripts/setup_pinecone.py index b53d7bdc..4f3b81fe 100644 --- a/examples/agent/langgraph-fsi-agent/before/scripts/setup_pinecone.py +++ b/examples/agent/langgraph-fsi-agent/before/scripts/setup_pinecone.py @@ -101,9 +101,7 @@ def upload_to_pinecone(chunked_docs, index_name: str, force_upload: bool = False # Create vector store and upload print(f"Uploading {len(chunked_docs)} chunks to Pinecone...") - vector_store = PineconeVectorStore.from_documents( - documents=chunked_docs, embedding=EMBEDDINGS, index_name=index_name - ) + vector_store = PineconeVectorStore.from_documents(documents=chunked_docs, embedding=EMBEDDINGS, index_name=index_name) print(f"Successfully uploaded {len(chunked_docs)} document chunks to Pinecone") return vector_store diff --git a/examples/agent/langgraph-open-telemetry/main.py b/examples/agent/langgraph-open-telemetry/main.py index 86a499cd..6d006482 100644 --- a/examples/agent/langgraph-open-telemetry/main.py +++ b/examples/agent/langgraph-open-telemetry/main.py @@ -4,6 +4,7 @@ import dotenv import openai +from splunk_ao import otel # LangGraph imports - this is what we're actually instrumenting from langgraph.graph import END, StateGraph @@ -13,8 +14,6 @@ from opentelemetry.sdk import trace as trace_sdk # SDK for creating traces from opentelemetry.sdk.resources import Resource -from splunk_ao import otel - dotenv.load_dotenv() logging.basicConfig(level=logging.DEBUG) @@ -85,7 +84,10 @@ def generate_response(state: AgentState): # Make the OpenAI API call - OpenAI instrumentation handles tracing response = client.chat.completions.create( - model="gpt-3.5-turbo", messages=[{"role": "user", "content": user_input}], max_tokens=300, temperature=0.7 + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": user_input}], + max_tokens=300, + temperature=0.7, ) # Extract the response content @@ -100,7 +102,7 @@ def generate_response(state: AgentState): except Exception as e: print(f"❌ Error calling OpenAI: {e}") - return {"llm_response": f"Error: {e!s}"} + return {"llm_response": f"Error: {str(e)}"} # Node 3: Format Answer diff --git a/examples/agent/langgraph-otel/main.py b/examples/agent/langgraph-otel/main.py index 358545b8..0cbd85a9 100644 --- a/examples/agent/langgraph-otel/main.py +++ b/examples/agent/langgraph-otel/main.py @@ -14,20 +14,25 @@ # "instrument" your code so you can see exactly what's happening during execution. # Core OpenTelemetry imports -# OpenAI imports for LLM integration -import openai - -# LangGraph imports - this is what we're actually instrumenting -from langgraph.graph import END, StateGraph +from opentelemetry.sdk import trace as trace_sdk # SDK for creating traces +from opentelemetry import trace as trace_api # API for interacting with traces +from opentelemetry.sdk.trace.export import ( + BatchSpanProcessor, +) # Efficiently batches spans before export +from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter, +) # Sends traces via HTTP # OpenInference is a specialized instrumentation library that understands AI frameworks # It automatically creates meaningful spans for LangChain/LangGraph operations from openinference.instrumentation.langchain import LangChainInstrumentor from openinference.instrumentation.openai import OpenAIInstrumentor -from opentelemetry import trace as trace_api # API for interacting with traces -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter # Sends traces via HTTP -from opentelemetry.sdk import trace as trace_sdk # SDK for creating traces -from opentelemetry.sdk.trace.export import BatchSpanProcessor # Efficiently batches spans before export + +# LangGraph imports - this is what we're actually instrumenting +from langgraph.graph import StateGraph, END + +# OpenAI imports for LLM integration +import openai # ============================================================================ # STEP 1: CONFIGURE API AUTHENTICATION @@ -163,7 +168,10 @@ def generate_response(state: AgentState): # Make the OpenAI API call - OpenAI instrumentation handles tracing response = client.chat.completions.create( - model="gpt-3.5-turbo", messages=[{"role": "user", "content": user_input}], max_tokens=300, temperature=0.7 + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": user_input}], + max_tokens=300, + temperature=0.7, ) # Extract the response content @@ -175,7 +183,7 @@ def generate_response(state: AgentState): except Exception as e: print(f"❌ Error calling OpenAI: {e}") - return {"llm_response": f"Error: {e!s}"} + return {"llm_response": f"Error: {str(e)}"} # Node 3: Format Answer diff --git a/examples/agent/langgraph-telecom-agent/app.py b/examples/agent/langgraph-telecom-agent/app.py index 991446ac..f54c6c55 100644 --- a/examples/agent/langgraph-telecom-agent/app.py +++ b/examples/agent/langgraph-telecom-agent/app.py @@ -102,9 +102,7 @@ async def main(msg: cl.Message) -> None: main_step.input = msg.content # Call the graph with the user's message and stream the response back to the user - async for response_msg in supervisor_agent.astream( - input=messages, stream_mode="updates", config=runnable_config - ): + async for response_msg in supervisor_agent.astream(input=messages, stream_mode="updates", config=runnable_config): # Debug: Log the response structure print(f"Response keys: {response_msg.keys()}") diff --git a/examples/agent/langgraph-telecom-agent/scripts/setup_pinecone.py b/examples/agent/langgraph-telecom-agent/scripts/setup_pinecone.py index bcc853bd..d8644451 100644 --- a/examples/agent/langgraph-telecom-agent/scripts/setup_pinecone.py +++ b/examples/agent/langgraph-telecom-agent/scripts/setup_pinecone.py @@ -106,9 +106,7 @@ def upload_to_pinecone(chunked_docs, index_name: str, force_upload: bool = False # Create vector store and upload print(f"Uploading {len(chunked_docs)} chunks to Pinecone...") - vector_store = PineconeVectorStore.from_documents( - documents=chunked_docs, embedding=EMBEDDINGS, index_name=index_name - ) + vector_store = PineconeVectorStore.from_documents(documents=chunked_docs, embedding=EMBEDDINGS, index_name=index_name) print(f"Successfully uploaded {len(chunked_docs)} document chunks to Pinecone") return vector_store diff --git a/examples/agent/langgraph-telecom-agent/src/tools/billing_tool.py b/examples/agent/langgraph-telecom-agent/src/tools/billing_tool.py index aba5b087..db644b02 100644 --- a/examples/agent/langgraph-telecom-agent/src/tools/billing_tool.py +++ b/examples/agent/langgraph-telecom-agent/src/tools/billing_tool.py @@ -36,18 +36,18 @@ def _run(self, customer_id: Optional[str] = None, query_type: str = "summary") - if query_type == "usage": return f""" -Usage Summary for {customer["name"]}: -- Data: {customer["data_used"]:.1f} GB used ({customer["data_limit"]}) +Usage Summary for {customer['name']}: +- Data: {customer['data_used']:.1f} GB used ({customer['data_limit']}) - Minutes: {random.randint(300, 800)} (Unlimited) - Texts: {random.randint(500, 2000)} (Unlimited) -- Average daily: {customer["data_used"] / 15:.2f} GB +- Average daily: {customer['data_used'] / 15:.2f} GB """ elif query_type == "plan": return f""" -Current Plan: {customer["plan"]} -- Monthly Cost: ${customer["monthly_charge"]:.2f} -- Data: {customer["data_limit"]} +Current Plan: {customer['plan']} +- Monthly Cost: ${customer['monthly_charge']:.2f} +- Data: {customer['data_limit']} - Talk & Text: Unlimited - 5G Access: Included @@ -72,10 +72,10 @@ def _run(self, customer_id: Optional[str] = None, query_type: str = "summary") - # Default summary return f""" -Account Summary for {customer["name"]}: -- Account: {customer["account"]} -- Plan: {customer["plan"]} -- Amount Due: ${customer["monthly_charge"]:.2f} -- Due Date: {customer["due_date"]} -- Data Used: {customer["data_used"]:.1f} GB ({customer["data_limit"]}) +Account Summary for {customer['name']}: +- Account: {customer['account']} +- Plan: {customer['plan']} +- Amount Due: ${customer['monthly_charge']:.2f} +- Due Date: {customer['due_date']} +- Data Used: {customer['data_used']:.1f} GB ({customer['data_limit']}) """ diff --git a/examples/agent/langgraph-telecom-agent/src/tools/technical_support_tool.py b/examples/agent/langgraph-telecom-agent/src/tools/technical_support_tool.py index 653a5400..0fd7184e 100644 --- a/examples/agent/langgraph-telecom-agent/src/tools/technical_support_tool.py +++ b/examples/agent/langgraph-telecom-agent/src/tools/technical_support_tool.py @@ -70,7 +70,7 @@ def _run(self, issue_type: Optional[str] = None, action: str = "help") -> str: Support Ticket Created: {ticket_id} Priority: High Response Time: Within 2 hours -Callback: {(datetime.now() + timedelta(hours=2)).strftime("%Y-%m-%d %H:%M")} +Callback: {(datetime.now() + timedelta(hours=2)).strftime('%Y-%m-%d %H:%M')} 24/7 Support: 1-800-TELECOM """ diff --git a/examples/agent/langgraph-traceloop/agent.py b/examples/agent/langgraph-traceloop/agent.py index cfaa8790..44e7b052 100644 --- a/examples/agent/langgraph-traceloop/agent.py +++ b/examples/agent/langgraph-traceloop/agent.py @@ -1,13 +1,13 @@ import os -from typing import Annotated, TypedDict - +from typing import TypedDict, Annotated import dotenv + import openai -from langchain_core.messages import BaseMessage -from langchain_core.tools import tool -from langchain_openai import ChatOpenAI from langgraph.graph import END, StateGraph, add_messages from langgraph.prebuilt import ToolNode +from langchain_core.tools import tool +from langchain_core.messages import BaseMessage +from langchain_openai import ChatOpenAI dotenv.load_dotenv() @@ -60,7 +60,10 @@ def generate_response_tool(user_input: str) -> str: # Make the OpenAI API call - Traceloop automatically traces this response = client.chat.completions.create( - model="gpt-3.5-turbo", messages=[{"role": "user", "content": user_input}], max_tokens=300, temperature=0.7 + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": user_input}], + max_tokens=300, + temperature=0.7, ) # Extract the response content @@ -69,13 +72,14 @@ def generate_response_tool(user_input: str) -> str: if not llm_response: print("No response from OpenAI") return "Error: No response from OpenAI" - print(f"Received response: '{llm_response[:100]}...'") + else: + print(f"Received response: '{llm_response[:100]}...'") return llm_response except Exception as e: print(f"Error calling OpenAI: {e}") - return f"Error: {e!s}" + return f"Error: {str(e)}" @tool @@ -144,7 +148,7 @@ def agent_node(state: AgentState): return {"messages": [response]} -def should_continue(state: AgentState) -> str: +def should_continue(state: AgentState): """ Determines whether to continue with tool calls or end. """ @@ -179,10 +183,19 @@ def create_agent(): workflow.set_entry_point("agent") # Add conditional edges - workflow.add_conditional_edges("agent", should_continue, {"tools": "tools", "end": END}) + workflow.add_conditional_edges( + "agent", + should_continue, + { + "tools": "tools", + "end": END, + }, + ) # After tools are called, go back to the agent workflow.add_edge("tools", "agent") # Compile the workflow - return workflow.compile() + app = workflow.compile() + + return app diff --git a/examples/agent/langgraph-traceloop/main.py b/examples/agent/langgraph-traceloop/main.py index 72102a3a..81b6ef44 100644 --- a/examples/agent/langgraph-traceloop/main.py +++ b/examples/agent/langgraph-traceloop/main.py @@ -1,14 +1,17 @@ import dotenv from agent import create_agent -from langchain_core.messages import HumanMessage from traceloop.sdk import Traceloop +from langchain_core.messages import HumanMessage dotenv.load_dotenv() Traceloop.init( app_name="LangGraph-Traceloop-Demo", disable_batch=False, # Enable batching for better performance - resource_attributes={"service.version": "1.0.0", "deployment.environment": "development"}, + resource_attributes={ + "service.version": "1.0.0", + "deployment.environment": "development", + }, ) @@ -45,7 +48,7 @@ messages = result.get("messages", []) for i, msg in enumerate(messages): msg_type = type(msg).__name__ - print(f"\n--- Message {i + 1} ({msg_type}) ---") + print(f"\n--- Message {i+1} ({msg_type}) ---") if hasattr(msg, "content") and msg.content: print(f"Content: {msg.content[:200]}...") diff --git a/examples/agent/microsoft-agent-framework/agent.py b/examples/agent/microsoft-agent-framework/agent.py index 48d9ce35..196c1b33 100644 --- a/examples/agent/microsoft-agent-framework/agent.py +++ b/examples/agent/microsoft-agent-framework/agent.py @@ -3,12 +3,11 @@ from agent_framework import openai, tool from agent_framework.observability import enable_instrumentation +from splunk_ao.otel import SplunkAOSpanProcessor, add_splunk_ao_span_processor from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from pydantic import Field -from splunk_ao.otel import SplunkAOSpanProcessor, add_splunk_ao_span_processor - # Set up the OTel tracer provider with the Splunk AO span processor tracer_provider = TracerProvider() galileo_processor = SplunkAOSpanProcessor() @@ -27,7 +26,9 @@ @tool(approval_mode="never_require") -def get_weather(location: Annotated[str, Field(description="The location to get the weather for.")]) -> str: +def get_weather( + location: Annotated[str, Field(description="The location to get the weather for.")], +) -> str: """Get the weather for a given location.""" conditions = ["sunny", "cloudy", "rainy", "stormy"] return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}C." @@ -42,7 +43,7 @@ def get_weather(location: Annotated[str, Field(description="The location to get ) -async def main() -> None: +async def main(): result = await agent.run("What's the weather like in Seattle?") print(result) diff --git a/examples/agent/minimal-agent-example/app.py b/examples/agent/minimal-agent-example/app.py index 9d388512..7fe10462 100644 --- a/examples/agent/minimal-agent-example/app.py +++ b/examples/agent/minimal-agent-example/app.py @@ -1,12 +1,14 @@ -import json import os -from collections.abc import Callable - import streamlit as st -from dotenv import load_dotenv +from splunk_ao import ( + log, + splunk_ao_context, + openai, +) from pydantic import BaseModel - -from splunk_ao import log, openai, splunk_ao_context +import json +from typing import Callable +from dotenv import load_dotenv load_dotenv() @@ -40,9 +42,14 @@ class BudgetRequest(BaseModel): "parameters": { "type": "object", "properties": { - "destination": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"} + "destination": { + "type": "string", + "description": "City and country e.g. Bogotá, Colombia", + }, }, - "required": ["destination"], + "required": [ + "destination", + ], "additionalProperties": False, }, "strict": True, @@ -56,8 +63,14 @@ class BudgetRequest(BaseModel): "parameters": { "type": "object", "properties": { - "destination": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"}, - "days": {"type": "integer", "description": "Number of days for the itinerary"}, + "destination": { + "type": "string", + "description": "City and country e.g. Bogotá, Colombia", + }, + "days": { + "type": "integer", + "description": "Number of days for the itinerary", + }, }, "required": ["destination", "days"], "additionalProperties": False, @@ -73,9 +86,14 @@ class BudgetRequest(BaseModel): "parameters": { "type": "object", "properties": { - "destination": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"} + "destination": { + "type": "string", + "description": "City and country e.g. Bogotá, Colombia", + }, }, - "required": ["destination"], + "required": [ + "destination", + ], "additionalProperties": False, }, "strict": True, @@ -89,8 +107,14 @@ class BudgetRequest(BaseModel): "parameters": { "type": "object", "properties": { - "destination": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"}, - "days": {"type": "integer", "description": "Number of days for the itinerary"}, + "destination": { + "type": "string", + "description": "City and country e.g. Bogotá, Colombia", + }, + "days": { + "type": "integer", + "description": "Number of days for the itinerary", + }, }, "required": ["destination", "days"], "additionalProperties": False, @@ -114,12 +138,12 @@ def get_weather_forecast(destination: str) -> str: # LLM call: Generate a destination overview. # ============================================================================= def generate_destination_overview(destination: str) -> str: - prompt = ( - f"Provide a brief overview of {destination}, including its top attractions, " - "cultural highlights, and essential travel tips." - ) + prompt = f"Provide a brief overview of {destination}, including its top attractions, " "cultural highlights, and essential travel tips." # Call the OpenAI API (assuming proper API key configuration) - response = client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": prompt}]) + response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": prompt}], + ) # Extract and return the text from the response. return response.choices[0].message.content.strip() @@ -136,7 +160,10 @@ def generate_itinerary(destination: str, days: int) -> str: - Every plan must contain a destination overview and an itinerary. - Only include a travel budget and weather info if the user requests it. """ - response = client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": prompt}]) + response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": prompt}], + ) return response.choices[0].message.content.strip() @@ -150,7 +177,10 @@ def estimate_travel_budget(destination: str, days: int, itinerary: str | None = "food, transportation, and activities. Provide a rough breakdown of the costs (in USD)." f"{itinerary_prompt}" ) - response = client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": prompt}]) + response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": prompt}], + ) return response.choices[0].message.content.strip() @@ -174,7 +204,9 @@ def assemble_travel_plan(query: str, info_callback: Callable[[str], None]) -> st """ response = client.chat.completions.create( - model="gpt-4-turbo", messages=[{"role": "user", "content": function_calling_prompt}], tools=tools + model="gpt-4-turbo", + messages=[{"role": "user", "content": function_calling_prompt}], + tools=tools, ) # Extract function call responses @@ -200,9 +232,7 @@ def assemble_travel_plan(query: str, info_callback: Callable[[str], None]) -> st if function_name == "generate_destination_overview": destination_overview_request = DestinationOverviewRequest(**function_args) info_callback(f"Generating a destination overview for {destination_overview_request.destination}...") - destination_overview = ( - destination_overview + "\n" + generate_destination_overview(destination_overview_request.destination) - ) + destination_overview = destination_overview + "\n" + generate_destination_overview(destination_overview_request.destination) elif function_name == "generate_itinerary": itinerary_request = ItineraryRequest(**function_args) @@ -216,7 +246,9 @@ def assemble_travel_plan(query: str, info_callback: Callable[[str], None]) -> st budget + "\n" + estimate_travel_budget( - destination=budget_request.destination, days=budget_request.days, itinerary=itinerary + destination=budget_request.destination, + days=budget_request.days, + itinerary=itinerary, ) ) @@ -247,14 +279,17 @@ def assemble_travel_plan(query: str, info_callback: Callable[[str], None]) -> st + "\n\nPlease package the information above into a plan that I can use for my next trip." ) - response = client.chat.completions.create(model="gpt-4o", messages=[{"role": "user", "content": assembly_prompt}]) + response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "user", "content": assembly_prompt}], + ) return response.choices[0].message.content.strip() # ============================================================================= # Main Streamlit App # ============================================================================= -def main() -> None: +def main(): st.title("Travel Itinerary Planner") st.write("Plan your next adventure with our AI-powered travel planner.") @@ -273,7 +308,7 @@ def main() -> None: # Create an empty container for results so that previous outputs are cleared on re-run. result_container = st.empty() - def info_callback(message) -> None: + def info_callback(message): info_placeholder.info(message) if plan_trip_clicked: diff --git a/examples/agent/minimal-agent-example/minimal_agent.py b/examples/agent/minimal-agent-example/minimal_agent.py index ec3da41b..5a824dfa 100644 --- a/examples/agent/minimal-agent-example/minimal_agent.py +++ b/examples/agent/minimal-agent-example/minimal_agent.py @@ -1,14 +1,11 @@ -import json import os +import json import warnings - -import openai -import questionary +from splunk_ao import log, splunk_ao_context, openai as splunk_ao_openai from dotenv import load_dotenv from rich.console import Console - -from splunk_ao import log, splunk_ao_context -from splunk_ao import openai as splunk_ao_openai +import questionary +import openai # Suppress Pydantic serializer warnings warnings.filterwarnings("ignore", message="Pydantic serializer warnings") @@ -47,7 +44,7 @@ def convert_text_to_arithmetic_expression(text): # Tool: Calculator for arithmetic operations @log(span_type="tool", name="calculate") -def calculate(expression) -> str | None: +def calculate(expression): """Perform a calculation based on the given expression.""" console.print(f"Calculating: {expression}") @@ -56,12 +53,12 @@ def calculate(expression) -> str | None: console.print(f"Result: {result}") return f"The result of {expression} is {result}" except Exception as e: - return f"Error calculating {expression}: {e!s}" + return f"Error calculating {expression}: {str(e)}" # Load tools from tools.json def get_tools(): - with open(os.path.join(os.path.dirname(__file__), "tools.json")) as f: + with open(os.path.join(os.path.dirname(__file__), "tools.json"), "r") as f: return json.load(f) @@ -86,7 +83,11 @@ def process_query(query): # Agent loop - continue until the LLM decides we're done while True: # Get next tool call from LLM - response = client.chat.completions.create(model="gpt-4o", messages=messages, tools=tools) + response = client.chat.completions.create( + model="gpt-4o", + messages=messages, + tools=tools, + ) # Convert the assistant message to a serializable dictionary assistant_message = response.choices[0].message @@ -100,7 +101,10 @@ def process_query(query): { "id": tool_call.id, "type": "function", - "function": {"name": tool_call.function.name, "arguments": tool_call.function.arguments}, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments, + }, } ) assistant_dict["tool_calls"] = tool_calls_list @@ -142,10 +146,15 @@ def process_query(query): results.append(result) # Create a summary - return "\n".join(results) if results else "No results produced." + if results: + summary = "\n".join(results) + else: + summary = "No results produced." + + return summary -def main() -> None: +def main(): console.print("[bold]Minimal Number Converter & Calculator[/bold]") query = questionary.text("Enter your query:", default="What's 4 + seven?").ask() diff --git a/examples/agent/pydantic-ai-support-agent/src/main.py b/examples/agent/pydantic-ai-support-agent/src/main.py index d47b7ac8..16da77bc 100644 --- a/examples/agent/pydantic-ai-support-agent/src/main.py +++ b/examples/agent/pydantic-ai-support-agent/src/main.py @@ -2,19 +2,20 @@ from dataclasses import dataclass from datetime import datetime +from typing import Optional +from splunk_ao.otel import SplunkAOSpanProcessor, add_splunk_ao_span_processor from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from pydantic import BaseModel from pydantic_ai import Agent, RunContext -from splunk_ao.otel import SplunkAOSpanProcessor, add_splunk_ao_span_processor - # Set up Splunk AO observability provider = TracerProvider() trace.set_tracer_provider(provider) add_splunk_ao_span_processor( - tracer_provider=provider, processor=SplunkAOSpanProcessor(project="pydantic-ai-support", logstream="default") + tracer_provider=provider, + processor=SplunkAOSpanProcessor(project="pydantic-ai-support", logstream="default"), ) # Enable instrumentation on all PydanticAI agents @@ -52,15 +53,42 @@ class SupportTicket: # --- Mock Database --- CUSTOMERS_DB: dict[str, Customer] = { - "C001": Customer("C001", "Alice Johnson", "alice@example.com", "enterprise", datetime(2022, 1, 15)), + "C001": Customer( + "C001", + "Alice Johnson", + "alice@example.com", + "enterprise", + datetime(2022, 1, 15), + ), "C002": Customer("C002", "Bob Smith", "bob@example.com", "premium", datetime(2023, 6, 20)), "C003": Customer("C003", "Carol White", "carol@example.com", "standard", datetime(2024, 3, 10)), } ORDERS_DB: dict[str, Order] = { - "ORD-1001": Order("ORD-1001", "C001", "Enterprise License", 5000.00, "delivered", datetime(2024, 11, 1)), - "ORD-1002": Order("ORD-1002", "C001", "Support Package", 1200.00, "shipped", datetime(2024, 12, 15)), - "ORD-1003": Order("ORD-1003", "C002", "Premium Subscription", 299.99, "pending", datetime(2025, 1, 5)), + "ORD-1001": Order( + "ORD-1001", + "C001", + "Enterprise License", + 5000.00, + "delivered", + datetime(2024, 11, 1), + ), + "ORD-1002": Order( + "ORD-1002", + "C001", + "Support Package", + 1200.00, + "shipped", + datetime(2024, 12, 15), + ), + "ORD-1003": Order( + "ORD-1003", + "C002", + "Premium Subscription", + 299.99, + "pending", + datetime(2025, 1, 5), + ), } TICKETS_DB: dict[str, SupportTicket] = {} @@ -76,7 +104,7 @@ class SupportDeps: # --- Response Model --- class SupportResponse(BaseModel): message: str - ticket_id: str | None = None + ticket_id: Optional[str] = None refund_processed: bool = False @@ -100,7 +128,9 @@ async def get_customer_info(ctx: RunContext[SupportDeps]) -> str: customer = CUSTOMERS_DB.get(ctx.deps.customer_id) if not customer: return "Customer not found." - return f"Customer: {customer.name}\nEmail: {customer.email}\nTier: {customer.tier}\nAccount since: {customer.account_created.strftime('%Y-%m-%d')}" + return ( + f"Customer: {customer.name}\n" f"Email: {customer.email}\n" f"Tier: {customer.tier}\n" f"Account since: {customer.account_created.strftime('%Y-%m-%d')}" + ) @support_agent.tool @@ -154,7 +184,11 @@ async def create_support_ticket(ctx: RunContext[SupportDeps], subject: str, prio TICKET_COUNTER += 1 ticket_id = f"TKT-{TICKET_COUNTER}" ticket = SupportTicket( - id=ticket_id, customer_id=ctx.deps.customer_id, subject=subject, priority=priority, status="open" + id=ticket_id, + customer_id=ctx.deps.customer_id, + subject=subject, + priority=priority, + status="open", ) TICKETS_DB[ticket_id] = ticket return f"Support ticket created: {ticket_id} (Priority: {priority})" diff --git a/examples/agent/startup-simulator-3000/__init__.py b/examples/agent/startup-simulator-3000/__init__.py index 98e9c2f8..ee0f238e 100644 --- a/examples/agent/startup-simulator-3000/__init__.py +++ b/examples/agent/startup-simulator-3000/__init__.py @@ -1,3 +1,8 @@ from tools.startup_simulator import StartupSimulatorTool -__all__ = ["KeywordExtractorTool", "SimpleAgent", "StartupSimulatorTool", "TextAnalyzerTool"] +__all__ = [ + "SimpleAgent", + "TextAnalyzerTool", + "KeywordExtractorTool", + "StartupSimulatorTool", +] diff --git a/examples/agent/startup-simulator-3000/agent.py b/examples/agent/startup-simulator-3000/agent.py index 2fb6d9e3..402bd4e8 100644 --- a/examples/agent/startup-simulator-3000/agent.py +++ b/examples/agent/startup-simulator-3000/agent.py @@ -4,29 +4,28 @@ This agent generates startup pitches using different tools based on the selected mode. """ -import json import os +import json +from typing import Dict, Any, List, Optional from datetime import datetime -from typing import Any +from dotenv import load_dotenv +from splunk_ao import SplunkAOLogger +from splunk_ao.openai import openai from agent_framework.agent import Agent -from agent_framework.llm.models import LLMConfig +from agent_framework.models import VerbosityLevel, ToolSelectionHooks from agent_framework.llm.openai_provider import OpenAIProvider -from agent_framework.models import ToolSelectionHooks, VerbosityLevel +from agent_framework.llm.models import LLMConfig from agent_framework.utils.logging import ConsoleAgentLogger from agent_framework.utils.tool_hooks import LoggingToolSelectionHooks -from dotenv import load_dotenv -from tools.hackernews_tool import HackerNewsTool -from tools.keyword_extraction import KeywordExtractorTool -from tools.news_api_tool import NewsAPITool -from tools.serious_startup_simulator import SeriousStartupSimulatorTool # Import all available tools from tools.startup_simulator import StartupSimulatorTool +from tools.serious_startup_simulator import SeriousStartupSimulatorTool +from tools.hackernews_tool import HackerNewsTool +from tools.news_api_tool import NewsAPITool from tools.text_analysis import TextAnalyzerTool - -from splunk_ao import SplunkAOLogger -from splunk_ao.openai import openai +from tools.keyword_extraction import KeywordExtractorTool # Load environment variables load_dotenv() @@ -49,10 +48,10 @@ class SimpleAgent(Agent): def __init__( self, verbosity: VerbosityLevel = VerbosityLevel.LOW, - logger: ConsoleAgentLogger | None = None, - tool_selection_hooks: ToolSelectionHooks | None = None, - metadata: dict[str, Any] | None = None, - llm_provider: OpenAIProvider | None = None, + logger: Optional[ConsoleAgentLogger] = None, + tool_selection_hooks: Optional[ToolSelectionHooks] = None, + metadata: Optional[Dict[str, Any]] = None, + llm_provider: Optional[OpenAIProvider] = None, mode: str = "silly", ): # Create default LLM config if not provided @@ -65,8 +64,7 @@ def __init__( agent_id=f"startup-agent-{mode}", verbosity=verbosity, logger=logger or ConsoleAgentLogger(f"startup-agent-{mode}"), - tool_selection_hooks=tool_selection_hooks - or LoggingToolSelectionHooks(logger or ConsoleAgentLogger(f"startup-agent-{mode}")), + tool_selection_hooks=tool_selection_hooks or LoggingToolSelectionHooks(logger or ConsoleAgentLogger(f"startup-agent-{mode}")), metadata=metadata or {}, llm_provider=llm_provider, ) @@ -99,16 +97,23 @@ def _register_tools(self) -> None: self.tool_registry.register(metadata=TextAnalyzerTool.get_metadata(), implementation=TextAnalyzerTool) # Keyword extraction tool - finds important words/phrases in text - self.tool_registry.register(metadata=KeywordExtractorTool.get_metadata(), implementation=KeywordExtractorTool) + self.tool_registry.register( + metadata=KeywordExtractorTool.get_metadata(), + implementation=KeywordExtractorTool, + ) # Startup simulator tool - generates silly, creative startup pitches # Used in "silly" mode - self.tool_registry.register(metadata=StartupSimulatorTool.get_metadata(), implementation=StartupSimulatorTool) + self.tool_registry.register( + metadata=StartupSimulatorTool.get_metadata(), + implementation=StartupSimulatorTool, + ) # Serious startup simulator tool - generates professional business plans # Used in "serious" mode self.tool_registry.register( - metadata=SeriousStartupSimulatorTool.get_metadata(), implementation=SeriousStartupSimulatorTool + metadata=SeriousStartupSimulatorTool.get_metadata(), + implementation=SeriousStartupSimulatorTool, ) # HackerNews tool - fetches trending tech stories for inspiration @@ -119,7 +124,7 @@ def _register_tools(self) -> None: # Used in "serious" mode to get professional context self.tool_registry.register(metadata=NewsAPITool.get_metadata(), implementation=NewsAPITool) - async def _format_result(self, task: str, results: list[tuple[str, Any]], galileo_logger: SplunkAOLogger) -> str: + async def _format_result(self, task: str, results: List[tuple[str, Any]], galileo_logger: SplunkAOLogger) -> str: """ Format the final result from tool executions. @@ -339,7 +344,10 @@ async def _execute_hackernews_tool(self, limit: int = 3, galileo_logger: SplunkA raise e async def _execute_news_api_tool( - self, category: str = "business", limit: int = 5, galileo_logger: SplunkAOLogger = None + self, + category: str = "business", + limit: int = 5, + galileo_logger: SplunkAOLogger = None, ) -> str: """ Execute the NewsAPI tool to get business news for context. @@ -456,7 +464,10 @@ async def _execute_startup_simulator( if startup_tool_class: startup_tool = startup_tool_class() startup_result = await startup_tool.execute( - industry=industry, audience=audience, random_word=random_word, hn_context=hn_context + industry=industry, + audience=audience, + random_word=random_word, + hn_context=hn_context, ) # Add LLM span for tool completion @@ -467,11 +478,7 @@ async def _execute_startup_simulator( model="startup_simulator", num_input_tokens=len(industry) + len(audience) + len(random_word) + len(hn_context), num_output_tokens=len(startup_result), - total_tokens=len(industry) - + len(audience) - + len(random_word) - + len(hn_context) - + len(startup_result), + total_tokens=len(industry) + len(audience) + len(random_word) + len(hn_context) + len(startup_result), duration_ns=0, ) @@ -545,7 +552,10 @@ async def _execute_serious_startup_simulator( if startup_tool_class: startup_tool = startup_tool_class() startup_result = await startup_tool.execute( - industry=industry, audience=audience, random_word=random_word, news_context=news_context + industry=industry, + audience=audience, + random_word=random_word, + news_context=news_context, ) # Add LLM span for tool completion @@ -556,11 +566,7 @@ async def _execute_serious_startup_simulator( model="serious_startup_simulator", num_input_tokens=len(industry) + len(audience) + len(random_word) + len(news_context), num_output_tokens=len(startup_result), - total_tokens=len(industry) - + len(audience) - + len(random_word) - + len(news_context) - + len(startup_result), + total_tokens=len(industry) + len(audience) + len(random_word) + len(news_context) + len(startup_result), duration_ns=0, ) @@ -614,7 +620,11 @@ async def run(self, task: str, industry: str = "", audience: str = "", random_wo # Store parameters for tool execution # These will be passed to individual tools as needed - self.task_parameters = {"industry": industry, "audience": audience, "random_word": random_word} + self.task_parameters = { + "industry": industry, + "audience": audience, + "random_word": random_word, + } # Log workflow start as JSON for observability # This helps us understand the agent's decision-making process @@ -655,9 +665,7 @@ async def run(self, task: str, industry: str = "", audience: str = "", random_wo if self.mode == "serious": # Step 1: Get news context first print("🔍 Step 1: Fetching business news for context...") - news_context = await self._execute_news_api_tool( - category="business", limit=5, galileo_logger=galileo_logger - ) + news_context = await self._execute_news_api_tool(category="business", limit=5, galileo_logger=galileo_logger) results.append(("news_api", news_context)) # Step 2: Generate serious startup pitch using the news context diff --git a/examples/agent/startup-simulator-3000/agent_framework/__init__.py b/examples/agent/startup-simulator-3000/agent_framework/__init__.py index e66ad1c0..39e869af 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/__init__.py +++ b/examples/agent/startup-simulator-3000/agent_framework/__init__.py @@ -2,17 +2,17 @@ from .agent import Agent from .config import AgentConfiguration -from .exceptions import AgentError, ToolExecutionError, ToolNotFoundError from .models import AgentMetadata, VerbosityLevel +from .exceptions import AgentError, ToolNotFoundError, ToolExecutionError __version__ = "0.1.0" __all__ = [ "Agent", "AgentConfiguration", - "AgentError", "AgentMetadata", - "ToolExecutionError", - "ToolNotFoundError", "VerbosityLevel", + "AgentError", + "ToolNotFoundError", + "ToolExecutionError", ] diff --git a/examples/agent/startup-simulator-3000/agent_framework/agent.py b/examples/agent/startup-simulator-3000/agent_framework/agent.py index e91b0f84..9fa7cd2e 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/agent.py +++ b/examples/agent/startup-simulator-3000/agent_framework/agent.py @@ -1,24 +1,32 @@ from abc import ABC, abstractmethod -from datetime import datetime -from typing import Any +from typing import Any, Dict, List, Optional from uuid import uuid4 - +from datetime import datetime from splunk_ao import log # 🔍 Splunk AO import - this is the main Splunk AO logging library +from .utils.logging import AgentLogger +from .utils.tool_registry import ToolRegistry -from .exceptions import ToolExecutionError, ToolNotFoundError +from .models import ( + TaskExecution, + VerbosityLevel, + TaskAnalysis, + ToolContext, + ToolSelectionHooks, + AgentConfig, +) from .llm.base import LLMProvider from .llm.models import LLMMessage -from .models import AgentConfig, TaskAnalysis, TaskExecution, ToolContext, ToolSelectionHooks, VerbosityLevel + from .utils.formatting import ( + display_task_header, display_analysis, display_chain_of_thought, - display_error, display_execution_plan, + display_error, display_final_result, - display_task_header, ) -from .utils.logging import AgentLogger -from .utils.tool_registry import ToolRegistry + +from .exceptions import ToolNotFoundError, ToolExecutionError class Agent(ABC): @@ -27,25 +35,27 @@ class Agent(ABC): def __init__( self, *args, - agent_id: str | None = None, + agent_id: Optional[str] = None, verbosity: VerbosityLevel = VerbosityLevel.LOW, - logger: AgentLogger | None = None, - tool_selection_hooks: ToolSelectionHooks | None = None, - metadata: dict[str, Any] | None = None, - llm_provider: LLMProvider | None = None, + logger: Optional[AgentLogger] = None, + tool_selection_hooks: Optional[ToolSelectionHooks] = None, + metadata: Optional[Dict[str, Any]] = None, + llm_provider: Optional[LLMProvider] = None, **kwargs, ): self.agent_id = agent_id or str(uuid4()) self.config = AgentConfig( - verbosity=verbosity, tool_selection_hooks=tool_selection_hooks, metadata=metadata or {} + verbosity=verbosity, + tool_selection_hooks=tool_selection_hooks, + metadata=metadata or {}, ) self.llm_provider = llm_provider self.tool_registry = ToolRegistry() - self.current_task: TaskExecution | None = None - self.state: dict[str, Any] = {} - self.message_history: list[dict[str, Any]] = [] + self.current_task: Optional[TaskExecution] = None + self.state: Dict[str, Any] = {} + self.message_history: List[Dict[str, Any]] = [] self.logger = logger - self._current_plan: TaskAnalysis | None = None + self._current_plan: Optional[TaskAnalysis] = None def _setup_logger(self, logger: AgentLogger) -> None: """Create and set up the logger after tools are registered""" @@ -61,7 +71,7 @@ def log_message(self, message: str, level: VerbosityLevel = VerbosityLevel.LOW) if self.config.verbosity.value >= level.value: print(message) - def _create_tool_context(self, tool_name: str, inputs: dict[str, Any]) -> ToolContext: + def _create_tool_context(self, tool_name: str, inputs: Dict[str, Any]) -> ToolContext: """Create a context object for tool execution""" if not self.current_task: raise ValueError("No active task") @@ -87,8 +97,12 @@ def _create_tool_context(self, tool_name: str, inputs: dict[str, Any]) -> ToolCo # This means every tool call will be tracked in your Splunk AO dashboard @log(span_type="tool", name="tool_execution") async def call_tool( - self, tool_name: str, inputs: dict[str, Any], execution_reasoning: str, context: dict[str, Any] - ) -> dict[str, Any]: + self, + tool_name: str, + inputs: Dict[str, Any], + execution_reasoning: str, + context: Dict[str, Any], + ) -> Dict[str, Any]: """Execute a tool and log the call with selection reasoning""" tool = self.tool_registry.get_tool(tool_name) if not tool: @@ -128,7 +142,7 @@ async def call_tool( await tool.hooks.after_execution(tool_context, None, error=e) raise - async def _execute_tool(self, tool_name: str, inputs: dict[str, Any]) -> dict[str, Any]: + async def _execute_tool(self, tool_name: str, inputs: Dict[str, Any]) -> Dict[str, Any]: """Execute a tool with given inputs""" tool_impl = self.tool_registry.get_implementation(tool_name) if not tool_impl: @@ -145,7 +159,7 @@ async def _execute_tool(self, tool_name: str, inputs: dict[str, Any]) -> dict[st except Exception as e: raise ToolExecutionError(tool_name, e) - def _create_planning_prompt(self, task: str) -> list[LLMMessage]: + def _create_planning_prompt(self, task: str) -> List[LLMMessage]: """Create prompt for task planning""" tools_description = "\n".join( [ @@ -217,9 +231,7 @@ async def plan_task(self, task: str) -> TaskAnalysis: display_task_header(task) try: - plan: TaskAnalysis = await self.llm_provider.generate_structured( - messages, TaskAnalysis, self.llm_provider.config - ) + plan: TaskAnalysis = await self.llm_provider.generate_structured(messages, TaskAnalysis, self.llm_provider.config) # Log the planning response if self.logger: @@ -246,7 +258,11 @@ async def plan_task(self, task: str) -> TaskAnalysis: async def run(self, task: str) -> str: """Execute a task and return the result""" self.current_task = TaskExecution( - task_id=str(uuid4()), agent_id=self.agent_id, input=task, start_time=datetime.now(), steps=[] + task_id=str(uuid4()), + agent_id=self.agent_id, + input=task, + start_time=datetime.now(), + steps=[], ) if self.logger: @@ -284,7 +300,7 @@ async def run(self, task: str) -> str: self.current_task.status = "completed" self._current_plan = None # Clear the plan - async def _execute_step(self, step: dict[str, Any], task: str, plan: TaskAnalysis) -> Any: + async def _execute_step(self, step: Dict[str, Any], task: str, plan: TaskAnalysis) -> Any: """Execute a single step in the plan""" tool_name = step["tool"] if not self.tool_registry.get_tool(tool_name): @@ -301,18 +317,21 @@ async def _execute_step(self, step: dict[str, Any], task: str, plan: TaskAnalysi await hooks.after_selection(tool_context, tool_name, 1.0, [step["reasoning"]]) # Then execute the tool - return await self.call_tool( + result = await self.call_tool( tool_name=tool_name, inputs=inputs, execution_reasoning=step["reasoning"], context={"task": task, "plan": plan}, ) + return result + @abstractmethod - async def _format_result(self, task: str, results: list[tuple[str, dict[str, Any]]]) -> str: + async def _format_result(self, task: str, results: List[tuple[str, Dict[str, Any]]]) -> str: """Format the final result from tool executions""" + pass - async def _map_inputs_to_tool(self, tool_name: str, task: str, input_mapping: dict[str, str]) -> dict[str, Any]: + async def _map_inputs_to_tool(self, tool_name: str, task: str, input_mapping: Dict[str, str]) -> Dict[str, Any]: """Map inputs based on tool schema""" tool = self.tool_registry.get_tool(tool_name) if not tool: @@ -358,9 +377,9 @@ async def _map_inputs_to_tool(self, tool_name: str, task: str, input_mapping: di elif self.state.has_variable(input_name): mapped_inputs[input_name] = self.state.get_variable(input_name) elif input_schema.get("type") == "string": - if (input_name == "news_context" and hasattr(self, "context_data")) or ( - input_name == "hn_context" and hasattr(self, "context_data") - ): + if input_name == "news_context" and hasattr(self, "context_data"): + mapped_inputs[input_name] = getattr(self, "context_data", "") + elif input_name == "hn_context" and hasattr(self, "context_data"): mapped_inputs[input_name] = getattr(self, "context_data", "") elif hasattr(self, "task_parameters") and input_name in self.task_parameters: mapped_inputs[input_name] = self.task_parameters[input_name] diff --git a/examples/agent/startup-simulator-3000/agent_framework/config.py b/examples/agent/startup-simulator-3000/agent_framework/config.py index 84c228c1..884a071d 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/config.py +++ b/examples/agent/startup-simulator-3000/agent_framework/config.py @@ -1,35 +1,35 @@ import os -from dataclasses import dataclass, field -from typing import Any - +from typing import Optional, Any, Dict, List from dotenv import load_dotenv - -from .llm.models import LLMConfig +from dataclasses import dataclass, field from .models import VerbosityLevel +from .llm.models import LLMConfig class EnvironmentError(Exception): """Raised when required environment variables are missing""" + pass + @dataclass class AgentConfiguration: """Configuration for the agent framework""" llm_config: LLMConfig = field(default_factory=lambda: LLMConfig(model="gpt-4", temperature=0.1)) - api_keys: dict[str, str] = field(default_factory=dict) + api_keys: Dict[str, str] = field(default_factory=dict) verbosity: VerbosityLevel = field(default=VerbosityLevel.LOW) - metadata: dict[str, Any] = field(default_factory=dict) + metadata: Dict[str, Any] = field(default_factory=dict) enable_logging: bool = field(default=True) enable_tool_selection: bool = field(default=True) @staticmethod - def get_env(key: str, default: str | None = None) -> str | None: + def get_env(key: str, default: Optional[str] = None) -> Optional[str]: """Get an environment variable with an optional default""" return os.getenv(key, default) @classmethod - def from_env(cls, required_keys: list[str], optional_keys: dict[str, str] | None = None) -> "AgentConfiguration": + def from_env(cls, required_keys: List[str], optional_keys: Optional[Dict[str, str]] = None) -> "AgentConfiguration": """Create configuration from environment variables Args: @@ -43,9 +43,7 @@ def from_env(cls, required_keys: list[str], optional_keys: dict[str, str] | None for key_name in required_keys: env_value = os.getenv(f"{key_name.upper()}_API_KEY") if not env_value: - raise EnvironmentError( - f"{key_name.upper()}_API_KEY environment variable is required. Please set it in your .env file" - ) + raise EnvironmentError(f"{key_name.upper()}_API_KEY environment variable is required. " "Please set it in your .env file") api_keys[key_name] = env_value # Load optional API keys @@ -58,7 +56,8 @@ def from_env(cls, required_keys: list[str], optional_keys: dict[str, str] | None # Create configuration return cls( llm_config=LLMConfig( - model=os.getenv("LLM_MODEL", "gpt-4"), temperature=float(os.getenv("LLM_TEMPERATURE", "0.1")) + model=os.getenv("LLM_MODEL", "gpt-4"), + temperature=float(os.getenv("LLM_TEMPERATURE", "0.1")), ), api_keys=api_keys, verbosity=VerbosityLevel(os.getenv("VERBOSITY", "low")), @@ -74,7 +73,7 @@ def with_overrides(self, **overrides) -> "AgentConfiguration": return AgentConfiguration(**config_dict) @classmethod - def from_dict(cls, config_dict: dict[str, Any]) -> "AgentConfiguration": + def from_dict(cls, config_dict: Dict[str, Any]) -> "AgentConfiguration": """Create configuration from dictionary - mainly used for testing""" if "api_keys" not in config_dict or "openai" not in config_dict["api_keys"]: raise ValueError("OpenAI API key must be provided in api_keys dictionary") diff --git a/examples/agent/startup-simulator-3000/agent_framework/exceptions.py b/examples/agent/startup-simulator-3000/agent_framework/exceptions.py index 36ba73a9..4a6e9e0f 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/exceptions.py +++ b/examples/agent/startup-simulator-3000/agent_framework/exceptions.py @@ -1,14 +1,20 @@ class AgentError(Exception): """Base class for agent framework exceptions""" + pass + class ToolError(AgentError): """Base class for tool-related errors""" + pass + class ToolNotFoundError(ToolError): """Raised when a requested tool is not found""" + pass + class ToolExecutionError(ToolError): """Raised when a tool execution fails""" @@ -16,16 +22,22 @@ class ToolExecutionError(ToolError): def __init__(self, tool_name: str, original_error: Exception): self.tool_name = tool_name self.original_error = original_error - super().__init__(f"Tool {tool_name} execution failed: {original_error!s}") + super().__init__(f"Tool {tool_name} execution failed: {str(original_error)}") class ConfigurationError(AgentError): """Raised when there's a configuration problem""" + pass + class PlanningError(AgentError): """Raised when task planning fails""" + pass + class StateError(AgentError): """Raised when there's a state-related error""" + + pass diff --git a/examples/agent/startup-simulator-3000/agent_framework/factory.py b/examples/agent/startup-simulator-3000/agent_framework/factory.py index e5072799..19cb84ee 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/factory.py +++ b/examples/agent/startup-simulator-3000/agent_framework/factory.py @@ -1,8 +1,9 @@ -from .agent import Agent +from typing import Optional, Type from .config import AgentConfiguration from .llm.base import LLMProvider from .llm.openai_provider import OpenAIProvider from .utils.logging import ConsoleAgentLogger +from .agent import Agent class AgentFactory: @@ -10,8 +11,8 @@ class AgentFactory: def __init__(self, config: AgentConfiguration): self.config = config - self._llm_provider: LLMProvider | None = None - self._logger: ConsoleAgentLogger | None = None + self._llm_provider: Optional[LLMProvider] = None + self._logger: Optional[ConsoleAgentLogger] = None def get_llm_provider(self) -> LLMProvider: """Get or create LLM provider""" @@ -22,24 +23,21 @@ def get_llm_provider(self) -> LLMProvider: raise ValueError("No LLM provider configured") return self._llm_provider - def get_logger(self, agent_id: str) -> ConsoleAgentLogger | None: + def get_logger(self, agent_id: str) -> Optional[ConsoleAgentLogger]: """Get logger if enabled""" if self.config.enable_logging: return ConsoleAgentLogger(agent_id) return None - def create_agent(self, agent_class: type[Agent], agent_id: str | None = None, **kwargs) -> Agent: + def create_agent(self, agent_class: Type[Agent], agent_id: Optional[str] = None, **kwargs) -> Agent: """Create an agent instance with proper dependencies""" # Create core dependencies llm_provider = self.get_llm_provider() logger = self.get_logger(agent_id) if agent_id else None # Create agent with injected dependencies - return agent_class( - agent_id=agent_id, - llm_provider=llm_provider, - logger=logger, - verbosity=self.config.verbosity, - metadata=self.config.metadata, - **kwargs, + agent = agent_class( + agent_id=agent_id, llm_provider=llm_provider, logger=logger, verbosity=self.config.verbosity, metadata=self.config.metadata, **kwargs ) + + return agent diff --git a/examples/agent/startup-simulator-3000/agent_framework/llm/__init__.py b/examples/agent/startup-simulator-3000/agent_framework/llm/__init__.py index eb2533ed..ccf06cf0 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/llm/__init__.py +++ b/examples/agent/startup-simulator-3000/agent_framework/llm/__init__.py @@ -1,7 +1,7 @@ """LLM Provider Package""" from .base import LLMProvider -from .models import LLMConfig, LLMMessage, LLMResponse +from .models import LLMMessage, LLMResponse, LLMConfig from .openai_provider import OpenAIProvider -__all__ = ["LLMConfig", "LLMMessage", "LLMProvider", "LLMResponse", "OpenAIProvider"] +__all__ = ["LLMProvider", "LLMMessage", "LLMResponse", "LLMConfig", "OpenAIProvider"] diff --git a/examples/agent/startup-simulator-3000/agent_framework/llm/base.py b/examples/agent/startup-simulator-3000/agent_framework/llm/base.py index 75397c98..9571ec6c 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/llm/base.py +++ b/examples/agent/startup-simulator-3000/agent_framework/llm/base.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod +from typing import List, Optional -from .models import LLMConfig, LLMMessage, LLMResponse +from .models import LLMMessage, LLMResponse, LLMConfig class LLMProvider(ABC): @@ -10,9 +11,11 @@ def __init__(self, config: LLMConfig): self.config = config @abstractmethod - async def generate(self, messages: list[LLMMessage], config: LLMConfig | None = None) -> LLMResponse: + async def generate(self, messages: List[LLMMessage], config: Optional[LLMConfig] = None) -> LLMResponse: """Generate a response from the LLM""" + pass @abstractmethod - async def generate_stream(self, messages: list[LLMMessage], config: LLMConfig | None = None): + async def generate_stream(self, messages: List[LLMMessage], config: Optional[LLMConfig] = None): """Generate a streaming response from the LLM""" + pass diff --git a/examples/agent/startup-simulator-3000/agent_framework/llm/models.py b/examples/agent/startup-simulator-3000/agent_framework/llm/models.py index 74878122..346b65b4 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/llm/models.py +++ b/examples/agent/startup-simulator-3000/agent_framework/llm/models.py @@ -1,5 +1,4 @@ -from typing import Any - +from typing import Any, Dict, List, Optional from pydantic import BaseModel, Field @@ -8,16 +7,16 @@ class LLMMessage(BaseModel): role: str content: str - name: str | None = None + name: Optional[str] = None class LLMResponse(BaseModel): """Structured response from LLM""" content: str - raw_response: dict[str, Any] = Field(default_factory=dict) - finish_reason: str | None = None - usage: dict[str, int] | None = None + raw_response: Dict[str, Any] = Field(default_factory=dict) + finish_reason: Optional[str] = None + usage: Optional[Dict[str, int]] = None class LLMConfig(BaseModel): @@ -25,23 +24,23 @@ class LLMConfig(BaseModel): model: str temperature: float = 0.7 - max_tokens: int | None = None + max_tokens: Optional[int] = None top_p: float = 1.0 frequency_penalty: float = 0.0 presence_penalty: float = 0.0 - stop: list[str] | None = None - custom_settings: dict[str, Any] = Field(default_factory=dict) + stop: Optional[List[str]] = None + custom_settings: Dict[str, Any] = Field(default_factory=dict) class ToolSelectionOutput(BaseModel): """Output from tool selection""" - selected_tools: list[str] = Field(description="Names of the selected tools in order of execution") + selected_tools: List[str] = Field(description="Names of the selected tools in order of execution") confidence: float = Field(description="Confidence score for the tool selection (0-1)", ge=0.0, le=1.0) - reasoning_steps: list[str] = Field(description="List of reasoning steps that led to the tool selection") + reasoning_steps: List[str] = Field(description="List of reasoning steps that led to the tool selection") @classmethod - def model_json_schema(cls) -> dict[str, Any]: + def model_json_schema(cls) -> Dict[str, Any]: """Get JSON schema with example""" schema = super().model_json_schema() schema["examples"] = [ diff --git a/examples/agent/startup-simulator-3000/agent_framework/llm/openai_provider.py b/examples/agent/startup-simulator-3000/agent_framework/llm/openai_provider.py index 70f2ab6d..098fc4d7 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/llm/openai_provider.py +++ b/examples/agent/startup-simulator-3000/agent_framework/llm/openai_provider.py @@ -1,14 +1,11 @@ +from typing import Any, Dict, List, Optional, AsyncGenerator, Type, TypeVar +from splunk_ao.openai import openai # Use Splunk AO's OpenAI wrapper +from pydantic import BaseModel import os -from collections.abc import AsyncGenerator -from typing import Any, TypeVar - from dotenv import load_dotenv -from pydantic import BaseModel - -from splunk_ao.openai import openai # Use Splunk AO's OpenAI wrapper from .base import LLMProvider -from .models import LLMConfig, LLMMessage, LLMResponse +from .models import LLMMessage, LLMResponse, LLMConfig T = TypeVar("T", bound=BaseModel) @@ -16,7 +13,7 @@ class OpenAIProvider(LLMProvider): """OpenAI implementation of LLM provider with support for both regular and project-based keys""" - def __init__(self, config: LLMConfig, organization: str | None = None): + def __init__(self, config: LLMConfig, organization: Optional[str] = None): super().__init__(config) # Load environment variables @@ -52,13 +49,18 @@ def __init__(self, config: LLMConfig, organization: str | None = None): else: print("Initialized OpenAI client with regular API key") - def _prepare_messages(self, messages: list[LLMMessage]) -> list[dict[str, Any]]: + def _prepare_messages(self, messages: List[LLMMessage]) -> List[Dict[str, Any]]: """Convert internal message format to OpenAI format""" return [ - {"role": msg.role, "content": msg.content, **({"name": msg.name} if msg.name else {})} for msg in messages + { + "role": msg.role, + "content": msg.content, + **({"name": msg.name} if msg.name else {}), + } + for msg in messages ] - def _prepare_config(self, config: LLMConfig | None = None) -> dict[str, Any]: + def _prepare_config(self, config: Optional[LLMConfig] = None) -> Dict[str, Any]: """Prepare configuration for OpenAI API""" cfg = config or self.config return { @@ -72,7 +74,7 @@ def _prepare_config(self, config: LLMConfig | None = None) -> dict[str, Any]: **cfg.custom_settings, } - async def generate(self, messages: list[LLMMessage], config: LLMConfig | None = None) -> LLMResponse: + async def generate(self, messages: List[LLMMessage], config: Optional[LLMConfig] = None) -> LLMResponse: """Generate a response using OpenAI""" openai_messages = self._prepare_messages(messages) api_config = self._prepare_config(config) @@ -91,14 +93,10 @@ async def generate(self, messages: list[LLMMessage], config: LLMConfig | None = ) except Exception as e: if "401" in str(e) and self.is_project_key: - raise ValueError( - f"Project-based key authentication failed. Please check your OPENAI_PROJECT_ID and ensure the project exists. Error: {e}" - ) + raise ValueError(f"Project-based key authentication failed. Please check your OPENAI_PROJECT_ID and ensure the project exists. Error: {e}") raise e - async def generate_stream( - self, messages: list[LLMMessage], config: LLMConfig | None = None - ) -> AsyncGenerator[LLMResponse, None]: + async def generate_stream(self, messages: List[LLMMessage], config: Optional[LLMConfig] = None) -> AsyncGenerator[LLMResponse, None]: """Generate a streaming response using OpenAI""" openai_messages = self._prepare_messages(messages) api_config = self._prepare_config(config) @@ -115,13 +113,14 @@ async def generate_stream( ) except Exception as e: if "401" in str(e) and self.is_project_key: - raise ValueError( - f"Project-based key authentication failed. Please check your OPENAI_PROJECT_ID and ensure the project exists. Error: {e}" - ) + raise ValueError(f"Project-based key authentication failed. Please check your OPENAI_PROJECT_ID and ensure the project exists. Error: {e}") raise e async def generate_structured( - self, messages: list[LLMMessage], output_model: type[T], config: LLMConfig | None = None + self, + messages: List[LLMMessage], + output_model: Type[T], + config: Optional[LLMConfig] = None, ) -> T: """Generate a response with structured output using function calling""" openai_messages = self._prepare_messages(messages) @@ -150,7 +149,5 @@ async def generate_structured( raise ValueError(f"Failed to parse structured output: {e}") except Exception as e: if "401" in str(e) and self.is_project_key: - raise ValueError( - f"Project-based key authentication failed. Please check your OPENAI_PROJECT_ID and ensure the project exists. Error: {e}" - ) + raise ValueError(f"Project-based key authentication failed. Please check your OPENAI_PROJECT_ID and ensure the project exists. Error: {e}") raise e diff --git a/examples/agent/startup-simulator-3000/agent_framework/llm/tool_models.py b/examples/agent/startup-simulator-3000/agent_framework/llm/tool_models.py index c62b605d..110f9adb 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/llm/tool_models.py +++ b/examples/agent/startup-simulator-3000/agent_framework/llm/tool_models.py @@ -1,5 +1,4 @@ -from typing import Any - +from typing import Any, Dict, List from pydantic import BaseModel, Field @@ -8,13 +7,13 @@ class TextAnalysis(BaseModel): complexity_score: float = Field(ge=0.0, le=1.0) readability_level: str - main_topics: list[str] - key_points: list[str] + main_topics: List[str] + key_points: List[str] analysis_summary: str - language_metrics: dict[str, Any] + language_metrics: Dict[str, Any] @classmethod - def model_json_schema(cls) -> dict[str, Any]: + def model_json_schema(cls) -> Dict[str, Any]: """Get JSON schema with example""" schema = super().model_json_schema() schema["examples"] = [ @@ -28,7 +27,11 @@ def model_json_schema(cls) -> dict[str, Any]: "Addresses performance optimization", ], "analysis_summary": "Technical text focusing on advanced AI concepts", - "language_metrics": {"sentence_count": 15, "average_sentence_length": 20, "vocabulary_richness": 0.8}, + "language_metrics": { + "sentence_count": 15, + "average_sentence_length": 20, + "vocabulary_richness": 0.8, + }, } ] return schema @@ -37,21 +40,28 @@ def model_json_schema(cls) -> dict[str, Any]: class KeywordExtraction(BaseModel): """Structured output for keyword extraction""" - keywords: list[str] = Field(min_items=1) - importance_scores: dict[str, float] - categories: dict[str, list[str]] + keywords: List[str] = Field(min_items=1) + importance_scores: Dict[str, float] + categories: Dict[str, List[str]] extraction_confidence: float = Field(ge=0.0, le=1.0) context_relevance: str @classmethod - def model_json_schema(cls) -> dict[str, Any]: + def model_json_schema(cls) -> Dict[str, Any]: """Get JSON schema with example""" schema = super().model_json_schema() schema["examples"] = [ { "keywords": ["machine learning", "neural networks", "optimization"], - "importance_scores": {"machine learning": 0.9, "neural networks": 0.85, "optimization": 0.7}, - "categories": {"technical": ["machine learning", "neural networks"], "methodology": ["optimization"]}, + "importance_scores": { + "machine learning": 0.9, + "neural networks": 0.85, + "optimization": 0.7, + }, + "categories": { + "technical": ["machine learning", "neural networks"], + "methodology": ["optimization"], + }, "extraction_confidence": 0.85, "context_relevance": "High relevance to AI/ML domain", } diff --git a/examples/agent/startup-simulator-3000/agent_framework/models.py b/examples/agent/startup-simulator-3000/agent_framework/models.py index a5ec8394..f01230e1 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/models.py +++ b/examples/agent/startup-simulator-3000/agent_framework/models.py @@ -1,10 +1,8 @@ -from dataclasses import dataclass, field from datetime import datetime -from enum import StrEnum -from typing import Any - -from pydantic import BaseModel, ConfigDict, Field - +from typing import Any, Dict, List, Optional +from pydantic import BaseModel, Field, ConfigDict +from enum import Enum +from dataclasses import dataclass, field from .utils.hooks import ToolHooks, ToolSelectionHooks @@ -13,10 +11,10 @@ class ToolMetadata(BaseModel): name: str = Field(description="Unique identifier for the tool") description: str = Field(description="Human-readable description of what the tool does") - tags: list[str] = Field(description="Categories/capabilities of the tool") - input_schema: dict[str, Any] = Field(description="JSON schema for tool inputs") - output_schema: dict[str, Any] = Field(description="JSON schema for tool outputs") - examples: list[dict[str, Any]] | None = Field(default=None, description="Example uses of the tool") + tags: List[str] = Field(description="Categories/capabilities of the tool") + input_schema: Dict[str, Any] = Field(description="JSON schema for tool inputs") + output_schema: Dict[str, Any] = Field(description="JSON schema for tool outputs") + examples: Optional[List[Dict[str, Any]]] = Field(default=None, description="Example uses of the tool") class ToolError(BaseModel): @@ -30,16 +28,14 @@ class AgentMetadata(BaseModel): name: str = Field(description="Name of the agent") description: str = Field(description="What the agent does") - capabilities: list[str] = Field(description="High-level capabilities") - tools: list[ToolMetadata] = Field(description="Tools available to this agent") + capabilities: List[str] = Field(description="High-level capabilities") + tools: List[ToolMetadata] = Field(description="Tools available to this agent") version: str = Field(default="1.0.0", description="Version of the agent") - custom_attributes: dict[str, Any] = Field( - default_factory=dict, description="Additional custom metadata for the agent" - ) + custom_attributes: Dict[str, Any] = Field(default_factory=dict, description="Additional custom metadata for the agent") model_config = ConfigDict(arbitrary_types_allowed=True) -class VerbosityLevel(StrEnum): +class VerbosityLevel(str, Enum): """Controls how much information is displayed to the user""" NONE = "none" # Only show final results @@ -51,15 +47,11 @@ class TaskAnalysis(BaseModel): """Analysis of a task using chain of thought reasoning""" input_analysis: str = Field(description="Analysis of the input, identifying key requirements and constraints") - available_tools: list[str] = Field(description="List of tools available for the task") - tool_capabilities: dict[str, list[str]] = Field(description="Mapping of tools to their key capabilities") - execution_plan: list[dict[str, Any]] = Field( - description="Ordered list of steps to execute, each with tool and reasoning" - ) - requirements_coverage: dict[str, list[str]] = Field( - description="How the identified requirements are covered by the planned steps" - ) - chain_of_thought: list[str] = Field(description="Chain of thought reasoning that led to this plan") + available_tools: List[str] = Field(description="List of tools available for the task") + tool_capabilities: Dict[str, List[str]] = Field(description="Mapping of tools to their key capabilities") + execution_plan: List[Dict[str, Any]] = Field(description="Ordered list of steps to execute, each with tool and reasoning") + requirements_coverage: Dict[str, List[str]] = Field(description="How the identified requirements are covered by the planned steps") + chain_of_thought: List[str] = Field(description="Chain of thought reasoning that led to this plan") @dataclass @@ -70,17 +62,17 @@ def __init__( self, task: str, tool_name: str, - inputs: dict[str, Any], - available_tools: list[dict[str, Any]], - previous_tools: list[str], - previous_results: list[Any], - previous_errors: list[Any], - message_history: list[dict[str, Any]], + inputs: Dict[str, Any], + available_tools: List[Dict[str, Any]], + previous_tools: List[str], + previous_results: List[Any], + previous_errors: List[Any], + message_history: List[Dict[str, Any]], agent_id: str, task_id: str, start_time: datetime, - metadata: dict[str, Any], - plan: TaskAnalysis | None = None, + metadata: Dict[str, Any], + plan: Optional[TaskAnalysis] = None, ): self.task = task self.tool_name = tool_name @@ -103,66 +95,77 @@ class Tool: name: str = field(metadata={"description": "Unique identifier for the tool"}) description: str = field(metadata={"description": "Human-readable description of what the tool does"}) - tags: list[str] = field(metadata={"description": "Categories or labels for the tool's capabilities"}) - input_schema: dict[str, Any] = field(metadata={"description": "JSON Schema defining expected input parameters"}) - output_schema: dict[str, Any] = field(metadata={"description": "JSON Schema defining the tool's output structure"}) - hooks: ToolHooks | None = field( - default=None, metadata={"description": "Optional hooks for tool execution lifecycle"} + tags: List[str] = field(metadata={"description": "Categories or labels for the tool's capabilities"}) + input_schema: Dict[str, Any] = field(metadata={"description": "JSON Schema defining expected input parameters"}) + output_schema: Dict[str, Any] = field(metadata={"description": "JSON Schema defining the tool's output structure"}) + hooks: Optional[ToolHooks] = field( + default=None, + metadata={"description": "Optional hooks for tool execution lifecycle"}, ) class ToolSelectionCriteria(BaseModel): """Criteria used for selecting a tool""" - required_tags: list[str] = Field(default_factory=list, description="Tags that a tool must have to be considered") - preferred_tags: list[str] = Field( - default_factory=list, description="Tags that are desired but not required in a tool" + required_tags: List[str] = Field(default_factory=list, description="Tags that a tool must have to be considered") + preferred_tags: List[str] = Field( + default_factory=list, + description="Tags that are desired but not required in a tool", ) - context_requirements: dict[str, Any] = Field( - default_factory=dict, description="Specific contextual requirements that influence tool selection" + context_requirements: Dict[str, Any] = Field( + default_factory=dict, + description="Specific contextual requirements that influence tool selection", ) - custom_rules: dict[str, Any] = Field( - default_factory=dict, description="Additional rules or criteria for tool selection" + custom_rules: Dict[str, Any] = Field( + default_factory=dict, + description="Additional rules or criteria for tool selection", ) class ToolSelectionReasoning(BaseModel): """Record of the reasoning process for tool selection""" - context: dict[str, Any] = Field(description="Current context and state when the tool selection was made") - considered_tools: list[str] = Field(description="Names of all tools that were evaluated for selection") + context: Dict[str, Any] = Field(description="Current context and state when the tool selection was made") + considered_tools: List[str] = Field(description="Names of all tools that were evaluated for selection") selection_criteria: ToolSelectionCriteria = Field(description="Criteria used to evaluate and select the tool") - reasoning_steps: list[str] = Field(description="Detailed steps of the decision-making process") + reasoning_steps: List[str] = Field(description="Detailed steps of the decision-making process") selected_tool: str = Field(description="Name of the tool that was ultimately chosen") - confidence_score: float = Field(description="Confidence level in the tool selection (0.0 to 1.0)", ge=0.0, le=1.0) + confidence_score: float = Field( + description="Confidence level in the tool selection (0.0 to 1.0)", + ge=0.0, + le=1.0, + ) class ToolCall(BaseModel): """Record of a tool invocation""" tool_name: str = Field(description="Name of the tool that was called") - inputs: dict[str, Any] = Field(description="Parameters passed to the tool during execution") - outputs: dict[str, Any] | None = Field(default=None, description="Results returned by the tool execution") - selection_reasoning: ToolSelectionReasoning | None = Field( - default=None, description="Reasoning process that led to selecting this tool" - ) + inputs: Dict[str, Any] = Field(description="Parameters passed to the tool during execution") + outputs: Optional[Dict[str, Any]] = Field(default=None, description="Results returned by the tool execution") + selection_reasoning: Optional[ToolSelectionReasoning] = Field(default=None, description="Reasoning process that led to selecting this tool") execution_reasoning: str = Field(description="Explanation of why this tool was executed") - timestamp: datetime = Field(default_factory=datetime.utcnow, description="UTC timestamp when the tool was called") + timestamp: datetime = Field( + default_factory=datetime.utcnow, + description="UTC timestamp when the tool was called", + ) success: bool = Field(default=True, description="Whether the tool execution completed successfully") - error: str | None = Field(default=None, description="Error message if the tool execution failed") + error: Optional[str] = Field(default=None, description="Error message if the tool execution failed") class ExecutionStep(BaseModel): """Record of a single step in the agent's execution""" - step_type: str = Field( - description="Category or type of execution step (e.g., 'task_received', 'processing', 'completion')" - ) + step_type: str = Field(description="Category or type of execution step (e.g., 'task_received', 'processing', 'completion')") description: str = Field(description="Human-readable description of what happened in this step") - timestamp: datetime = Field(default_factory=datetime.utcnow, description="UTC timestamp when the step occurred") - tool_calls: list[ToolCall] = Field(default_factory=list, description="Tools that were called during this step") - intermediate_state: dict[str, Any] | None = Field( - default=None, description="State or context information captured during this step" + timestamp: datetime = Field( + default_factory=datetime.utcnow, + description="UTC timestamp when the step occurred", + ) + tool_calls: List[ToolCall] = Field(default_factory=list, description="Tools that were called during this step") + intermediate_state: Optional[Dict[str, Any]] = Field( + default=None, + description="State or context information captured during this step", ) @@ -172,16 +175,21 @@ class TaskExecution(BaseModel): task_id: str = Field(description="Unique identifier for this task execution") agent_id: str = Field(description="Identifier of the agent executing the task") input: str = Field(description="Original input or request given to the agent") - steps: list[ExecutionStep] = Field( - default_factory=list, description="Sequence of steps taken during task execution" + steps: List[ExecutionStep] = Field( + default_factory=list, + description="Sequence of steps taken during task execution", + ) + output: Optional[str] = Field(default=None, description="Final result or response from the task execution") + start_time: datetime = Field( + default_factory=datetime.utcnow, + description="UTC timestamp when task execution began", ) - output: str | None = Field(default=None, description="Final result or response from the task execution") - start_time: datetime = Field(default_factory=datetime.utcnow, description="UTC timestamp when task execution began") - end_time: datetime | None = Field(default=None, description="UTC timestamp when task execution completed") + end_time: Optional[datetime] = Field(default=None, description="UTC timestamp when task execution completed") status: str = Field( - default="in_progress", description="Current status of the task (e.g., 'in_progress', 'completed', 'failed')" + default="in_progress", + description="Current status of the task (e.g., 'in_progress', 'completed', 'failed')", ) - error: str | None = Field(default=None, description="Error message if the task execution failed") + error: Optional[str] = Field(default=None, description="Error message if the task execution failed") @dataclass @@ -189,15 +197,15 @@ class AgentConfig: """Configuration for an agent""" verbosity: VerbosityLevel = field( - default=VerbosityLevel.LOW, metadata={"description": "Level of detail to display to the user"} + default=VerbosityLevel.LOW, + metadata={"description": "Level of detail to display to the user"}, ) # logger: Optional[AgentLogger] = field( # default=None, # metadata={"description": "Optional logger for recording agent activity"} # ) - tool_selection_hooks: ToolSelectionHooks | None = field( - default=None, metadata={"description": "Hooks for tool selection lifecycle"} - ) - metadata: dict[str, Any] = field( - default_factory=dict, metadata={"description": "Additional configuration metadata"} + tool_selection_hooks: Optional[ToolSelectionHooks] = field(default=None, metadata={"description": "Hooks for tool selection lifecycle"}) + metadata: Dict[str, Any] = field( + default_factory=dict, + metadata={"description": "Additional configuration metadata"}, ) diff --git a/examples/agent/startup-simulator-3000/agent_framework/prompts/templates.py b/examples/agent/startup-simulator-3000/agent_framework/prompts/templates.py index b84f12a2..ddbeeb24 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/prompts/templates.py +++ b/examples/agent/startup-simulator-3000/agent_framework/prompts/templates.py @@ -1,6 +1,5 @@ +from typing import Dict, Any from pathlib import Path -from typing import Any - import jinja2 @@ -9,7 +8,9 @@ class PromptTemplate: def __init__(self, template_path: str): self.env = jinja2.Environment( - loader=jinja2.FileSystemLoader(Path(__file__).parent / "templates"), trim_blocks=True, lstrip_blocks=True + loader=jinja2.FileSystemLoader(Path(__file__).parent / "templates"), + trim_blocks=True, + lstrip_blocks=True, ) self.template = self.env.get_template(template_path) @@ -22,7 +23,7 @@ class PromptLibrary: """Central repository for prompt templates""" def __init__(self): - self.templates: dict[str, PromptTemplate] = {} + self.templates: Dict[str, PromptTemplate] = {} self._load_templates() def _load_templates(self) -> None: diff --git a/examples/agent/startup-simulator-3000/agent_framework/state.py b/examples/agent/startup-simulator-3000/agent_framework/state.py index 5dfcef09..28c222cb 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/state.py +++ b/examples/agent/startup-simulator-3000/agent_framework/state.py @@ -1,6 +1,6 @@ +from typing import Any, Dict, Optional from dataclasses import dataclass, field from datetime import datetime -from typing import Any @dataclass @@ -8,15 +8,15 @@ class AgentState: """Container for agent state management""" # General state - variables: dict[str, Any] = field(default_factory=dict) + variables: Dict[str, Any] = field(default_factory=dict) # Tool execution state - tool_results: dict[str, Any] = field(default_factory=dict) - last_tool: str | None = None + tool_results: Dict[str, Any] = field(default_factory=dict) + last_tool: Optional[str] = None # Task state - task_start_time: datetime | None = None - task_variables: dict[str, Any] = field(default_factory=dict) + task_start_time: Optional[datetime] = None + task_variables: Dict[str, Any] = field(default_factory=dict) def has_variable(self, name: str) -> bool: """Check if a variable exists in state""" diff --git a/examples/agent/startup-simulator-3000/agent_framework/tools/base.py b/examples/agent/startup-simulator-3000/agent_framework/tools/base.py index 7a37004e..49538d79 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/tools/base.py +++ b/examples/agent/startup-simulator-3000/agent_framework/tools/base.py @@ -1,6 +1,5 @@ from abc import ABC, abstractmethod -from typing import Any, ClassVar - +from typing import Any, Dict, ClassVar, Type from ..models import ToolMetadata @@ -8,10 +7,11 @@ class BaseTool(ABC): """Base class for all tools""" # Tool metadata as class variables - metadata: ClassVar[type[ToolMetadata]] + metadata: ClassVar[Type[ToolMetadata]] def __init__(self): """Initialize the base tool""" + pass @classmethod def get_metadata(cls) -> ToolMetadata: @@ -20,6 +20,6 @@ def get_metadata(cls) -> ToolMetadata: return cls.metadata() # This will use the default values defined in the metadata class @abstractmethod - async def execute(self, **inputs: Any) -> dict[str, Any]: + async def execute(self, **inputs: Any) -> Dict[str, Any]: """Execute the tool with given inputs""" raise NotImplementedError("Tool must implement execute method") diff --git a/examples/agent/startup-simulator-3000/agent_framework/tools/registry.py b/examples/agent/startup-simulator-3000/agent_framework/tools/registry.py index 70dcd263..adb9bb39 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/tools/registry.py +++ b/examples/agent/startup-simulator-3000/agent_framework/tools/registry.py @@ -1,5 +1,4 @@ -from typing import TypeVar - +from typing import Dict, List, Optional, Type, TypeVar from ..models import Tool, ToolMetadata from ..tools.base import BaseTool @@ -10,10 +9,10 @@ class ToolRegistry: """Central registry for tool management""" def __init__(self): - self.tools: dict[str, Tool] = {} - self._implementations: dict[str, type[BaseTool]] = {} + self.tools: Dict[str, Tool] = {} + self._implementations: Dict[str, Type["BaseTool"]] = {} - def register(self, *, metadata: T, implementation: type["BaseTool"]) -> None: + def register(self, *, metadata: T, implementation: Type["BaseTool"]) -> None: """Register a tool and its implementation""" if metadata.name in self.tools: raise ValueError(f"Tool {metadata.name} is already registered") @@ -31,18 +30,18 @@ def register(self, *, metadata: T, implementation: type["BaseTool"]) -> None: self.tools[metadata.name] = tool self._implementations[metadata.name] = implementation - def get_tool(self, name: str) -> Tool | None: + def get_tool(self, name: str) -> Optional[Tool]: """Get tool by name""" return self.tools.get(name) - def get_implementation(self, name: str) -> type["BaseTool"] | None: + def get_implementation(self, name: str) -> Optional[Type["BaseTool"]]: """Get tool implementation by name""" return self._implementations.get(name) - def list_tools(self) -> list[Tool]: + def list_tools(self) -> List[Tool]: """Get list of all registered tools""" return list(self.tools.values()) - def get_tools_by_tags(self, tags: list[str]) -> list[Tool]: + def get_tools_by_tags(self, tags: List[str]) -> List[Tool]: """Get tools that have all specified tags""" return [tool for tool in self.tools.values() if all(tag in tool.tags for tag in tags)] diff --git a/examples/agent/startup-simulator-3000/agent_framework/utils/formatting.py b/examples/agent/startup-simulator-3000/agent_framework/utils/formatting.py index fda5fbf6..c3c001e6 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/utils/formatting.py +++ b/examples/agent/startup-simulator-3000/agent_framework/utils/formatting.py @@ -1,13 +1,12 @@ """Utilities for formatting and displaying output""" import json -from typing import Any - +from typing import Any, Dict, List from rich.console import Console -from rich.markdown import Markdown from rich.panel import Panel -from rich.syntax import Syntax +from rich.markdown import Markdown from rich.table import Table +from rich.syntax import Syntax console = Console() @@ -17,17 +16,23 @@ def format_json(data: Any) -> str: return json.dumps(data, indent=2, default=str) -def display_task_header(task: str) -> None: +def display_task_header(task: str): """Display a task header""" - console.print(Panel(f"[bold blue]Task:[/bold blue] {task}", title="🤖 Agent Task", border_style="blue")) + console.print( + Panel( + f"[bold blue]Task:[/bold blue] {task}", + title="🤖 Agent Task", + border_style="blue", + ) + ) -def display_analysis(analysis: str) -> None: +def display_analysis(analysis: str): """Display task analysis""" console.print(Panel(Markdown(analysis), title="📋 Task Analysis", border_style="green")) -def display_chain_of_thought(steps: list[str]) -> None: +def display_chain_of_thought(steps: List[str]): """Display chain of thought reasoning""" table = Table(title="🤔 Chain of Thought", show_header=False, border_style="cyan") table.add_column("Step", style="dim") @@ -39,7 +44,7 @@ def display_chain_of_thought(steps: list[str]) -> None: console.print(table) -def display_execution_plan(plan: list[dict[str, Any]]) -> None: +def display_execution_plan(plan: List[Dict[str, Any]]): """Display execution plan""" table = Table(title="📝 Execution Plan", border_style="magenta") table.add_column("Tool", style="bold cyan") @@ -51,7 +56,7 @@ def display_execution_plan(plan: list[dict[str, Any]]) -> None: console.print(table) -def display_tool_result(tool_name: str, result: dict[str, Any]) -> None: +def display_tool_result(tool_name: str, result: Dict[str, Any]): """Display tool execution result""" if isinstance(result, (dict, list)): result_display = Syntax(format_json(result), "json", theme="monokai", word_wrap=True) @@ -61,11 +66,18 @@ def display_tool_result(tool_name: str, result: dict[str, Any]) -> None: console.print(Panel(result_display, title=f"🔧 {tool_name} Result", border_style="yellow")) -def display_final_result(result: str) -> None: +def display_final_result(result: str): """Display final combined result""" - console.print(Panel(Markdown(result), title="✨ Final Result", border_style="green", padding=(1, 2))) + console.print( + Panel( + Markdown(result), + title="✨ Final Result", + border_style="green", + padding=(1, 2), + ) + ) -def display_error(error: str) -> None: +def display_error(error: str): """Display error message""" console.print(Panel(f"[bold red]Error:[/bold red] {error}", title="❌ Error", border_style="red")) diff --git a/examples/agent/startup-simulator-3000/agent_framework/utils/hooks.py b/examples/agent/startup-simulator-3000/agent_framework/utils/hooks.py index d9185235..b5e37087 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/utils/hooks.py +++ b/examples/agent/startup-simulator-3000/agent_framework/utils/hooks.py @@ -1,7 +1,7 @@ -from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional from dataclasses import dataclass, field from datetime import datetime -from typing import Any +from abc import ABC, abstractmethod @dataclass @@ -10,19 +10,15 @@ class ToolContext: task: str = field(metadata={"description": "The current task being executed"}) tool_name: str = field(metadata={"description": "Name of the tool being executed"}) - inputs: dict[str, Any] = field(metadata={"description": "Input parameters for the tool"}) - previous_tools: list[str] = field( - metadata={"description": "List of tools that were previously executed in this task"} - ) - previous_results: list[dict[str, Any]] = field(metadata={"description": "Results from previous tool executions"}) - previous_errors: list[str] = field(metadata={"description": "Errors from previous tool executions"}) - message_history: list[dict[str, Any]] = field( - metadata={"description": "Complete history of messages and tool executions"} - ) + inputs: Dict[str, Any] = field(metadata={"description": "Input parameters for the tool"}) + previous_tools: List[str] = field(metadata={"description": "List of tools that were previously executed in this task"}) + previous_results: List[Dict[str, Any]] = field(metadata={"description": "Results from previous tool executions"}) + previous_errors: List[str] = field(metadata={"description": "Errors from previous tool executions"}) + message_history: List[Dict[str, Any]] = field(metadata={"description": "Complete history of messages and tool executions"}) agent_id: str = field(metadata={"description": "ID of the agent executing the tool"}) task_id: str = field(metadata={"description": "ID of the current task"}) start_time: datetime = field(metadata={"description": "When the task started"}) - metadata: dict[str, Any] = field(metadata={"description": "Additional metadata from agent configuration"}) + metadata: Dict[str, Any] = field(metadata={"description": "Additional metadata from agent configuration"}) class ToolHooks(ABC): @@ -31,10 +27,12 @@ class ToolHooks(ABC): @abstractmethod async def before_execution(self, context: ToolContext) -> None: """Called before tool execution with full context""" + pass @abstractmethod - async def after_execution(self, context: ToolContext, result: Any, error: Exception | None = None) -> None: + async def after_execution(self, context: ToolContext, result: Any, error: Optional[Exception] = None) -> None: """Called after tool execution with the result or error""" + pass class ToolSelectionHooks(ABC): @@ -42,6 +40,11 @@ class ToolSelectionHooks(ABC): @abstractmethod async def after_selection( - self, context: ToolContext, selected_tool: str, confidence: float, reasoning: list[str] + self, + context: ToolContext, + selected_tool: str, + confidence: float, + reasoning: List[str], ) -> None: """Called after tool selection with selection details""" + pass diff --git a/examples/agent/startup-simulator-3000/agent_framework/utils/logging.py b/examples/agent/startup-simulator-3000/agent_framework/utils/logging.py index 5c95cde0..aaf4fa9b 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/utils/logging.py +++ b/examples/agent/startup-simulator-3000/agent_framework/utils/logging.py @@ -1,12 +1,11 @@ -import json -from abc import ABC, abstractmethod +from typing import Any, Dict, List from datetime import datetime -from typing import Any - -from agent_framework.utils.hooks import ToolHooks, ToolSelectionHooks from rich.console import Console from rich.panel import Panel from rich.theme import Theme +import json +from abc import ABC, abstractmethod +from agent_framework.utils.hooks import ToolHooks, ToolSelectionHooks # Create a custom theme for our logger theme = Theme( @@ -35,7 +34,6 @@ def get_galileo_logger(): if _global_galileo_logger is None: import os - from dotenv import load_dotenv # Load environment variables @@ -79,38 +77,47 @@ def __init__(self, agent_id: str): @abstractmethod def info(self, message: str, **kwargs) -> None: """Log an informational message""" + pass @abstractmethod def warning(self, message: str, **kwargs) -> None: """Log a warning message""" + pass @abstractmethod def error(self, message: str, **kwargs) -> None: """Log an error message""" + pass @abstractmethod def debug(self, message: str, **kwargs) -> None: """Log a debug message""" + pass @abstractmethod - def _write_log(self, log_entry: dict[str, Any]) -> None: + def _write_log(self, log_entry: Dict[str, Any]) -> None: """Write a log entry""" + pass @abstractmethod def _sanitize_for_json(self, obj: Any) -> Any: """Sanitize an object for JSON serialization""" + pass @abstractmethod async def on_agent_planning(self, planning_prompt: str) -> None: """Log the agent planning prompt""" + pass @abstractmethod def on_agent_start(self, initial_task: str) -> None: """Log the agent execution prompt""" + pass @abstractmethod - async def on_agent_done(self, result: str, message_history: list[dict[str, Any]]) -> None: + async def on_agent_done(self, result: str, message_history: List[Dict[str, Any]]) -> None: """Log the agent completion""" + pass def get_tool_hooks(self) -> ToolHooks: """Get tool hooks for this logger""" @@ -148,17 +155,18 @@ def debug(self, message: str, **kwargs) -> None: if kwargs: console.print(Panel(json.dumps(kwargs, indent=2), title="Additional Info")) - def _write_log(self, log_entry: dict[str, Any]) -> None: + def _write_log(self, log_entry: Dict[str, Any]) -> None: pass # Console logger doesn't need to write to file def _sanitize_for_json(self, obj: Any) -> Any: if isinstance(obj, (str, int, float, bool, type(None))): return obj - if isinstance(obj, (list, tuple)): + elif isinstance(obj, (list, tuple)): return [self._sanitize_for_json(item) for item in obj] - if isinstance(obj, dict): + elif isinstance(obj, dict): return {str(k): self._sanitize_for_json(v) for k, v in obj.items()} - return str(obj) + else: + return str(obj) async def on_agent_planning(self, planning_prompt: str) -> None: self.info(f"Planning: {planning_prompt}") @@ -166,5 +174,5 @@ async def on_agent_planning(self, planning_prompt: str) -> None: def on_agent_start(self, initial_task: str) -> None: self.info(f"Starting task: {initial_task}") - async def on_agent_done(self, result: str, message_history: list[dict[str, Any]]) -> None: + async def on_agent_done(self, result: str, message_history: List[Dict[str, Any]]) -> None: self.info(f"Task completed: {result}") diff --git a/examples/agent/startup-simulator-3000/agent_framework/utils/tool_hooks.py b/examples/agent/startup-simulator-3000/agent_framework/utils/tool_hooks.py index d45dd9ea..0e7f0400 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/utils/tool_hooks.py +++ b/examples/agent/startup-simulator-3000/agent_framework/utils/tool_hooks.py @@ -1,5 +1,4 @@ -from typing import Any - +from typing import Any, List, Optional from .hooks import ToolContext, ToolHooks, ToolSelectionHooks from .logging import AgentLogger @@ -11,13 +10,25 @@ def __init__(self, logger: AgentLogger): self.logger = logger async def before_execution(self, context: ToolContext) -> None: - self.logger.info(f"Executing tool: {context.tool_name}", inputs=context.inputs, task_id=context.task_id) + self.logger.info( + f"Executing tool: {context.tool_name}", + inputs=context.inputs, + task_id=context.task_id, + ) - async def after_execution(self, context: ToolContext, result: Any, error: Exception | None = None) -> None: + async def after_execution(self, context: ToolContext, result: Any, error: Optional[Exception] = None) -> None: if error: - self.logger.error(f"Tool execution failed: {context.tool_name}", error=str(error), task_id=context.task_id) + self.logger.error( + f"Tool execution failed: {context.tool_name}", + error=str(error), + task_id=context.task_id, + ) else: - self.logger.info(f"Tool execution completed: {context.tool_name}", result=result, task_id=context.task_id) + self.logger.info( + f"Tool execution completed: {context.tool_name}", + result=result, + task_id=context.task_id, + ) class LoggingToolSelectionHooks(ToolSelectionHooks): @@ -27,10 +38,17 @@ def __init__(self, logger: AgentLogger): self.logger = logger async def after_selection( - self, context: ToolContext, selected_tool: str, confidence: float, reasoning: list[str] + self, + context: ToolContext, + selected_tool: str, + confidence: float, + reasoning: List[str], ) -> None: self.logger.info( - f"Selected tool: {selected_tool}", confidence=confidence, reasoning=reasoning, task_id=context.task_id + f"Selected tool: {selected_tool}", + confidence=confidence, + reasoning=reasoning, + task_id=context.task_id, ) diff --git a/examples/agent/startup-simulator-3000/agent_framework/utils/tool_registry.py b/examples/agent/startup-simulator-3000/agent_framework/utils/tool_registry.py index e79e15af..e1b15690 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/utils/tool_registry.py +++ b/examples/agent/startup-simulator-3000/agent_framework/utils/tool_registry.py @@ -1,6 +1,5 @@ +from typing import Dict, List, Optional, Type, Any from dataclasses import dataclass, field -from typing import Any - from ..models import Tool, ToolMetadata from ..tools.base import BaseTool @@ -9,10 +8,10 @@ class ToolRegistry: """Central registry for tool management""" - tools: dict[str, Tool] = field(default_factory=dict) - _implementations: dict[str, type["BaseTool"]] = field(default_factory=dict) + tools: Dict[str, Tool] = field(default_factory=dict) + _implementations: Dict[str, Type["BaseTool"]] = field(default_factory=dict) - def register(self, *, metadata: ToolMetadata, implementation: type["BaseTool"]) -> None: + def register(self, *, metadata: ToolMetadata, implementation: Type["BaseTool"]) -> None: """Register a tool and its implementation""" if metadata.name in self.tools: raise ValueError(f"Tool {metadata.name} is already registered") @@ -30,27 +29,27 @@ def register(self, *, metadata: ToolMetadata, implementation: type["BaseTool"]) self.tools[metadata.name] = tool self._implementations[metadata.name] = implementation - def get_tool(self, name: str) -> Tool | None: + def get_tool(self, name: str) -> Optional[Tool]: """Get tool by name""" return self.tools.get(name) - def get_implementation(self, name: str) -> type["BaseTool"] | None: + def get_implementation(self, name: str) -> Optional[Type["BaseTool"]]: """Get tool implementation by name""" return self._implementations.get(name) - def list_tools(self) -> list[Tool]: + def list_tools(self) -> List[Tool]: """Get list of all registered tools""" return list(self.tools.values()) - def get_tools_by_tags(self, tags: list[str]) -> list[Tool]: + def get_tools_by_tags(self, tags: List[str]) -> List[Tool]: """Get tools that have all specified tags""" return [tool for tool in self.tools.values() if all(tag in tool.tags for tag in tags)] - def get_all_tools(self) -> dict[str, Tool]: + def get_all_tools(self) -> Dict[str, Tool]: """Get all registered tools""" return self.tools.copy() - def get_formatted_tools(self) -> list[dict[str, Any]]: + def get_formatted_tools(self) -> List[Dict[str, Any]]: """Format tools into OpenAI function calling format Returns a list of tools formatted as: diff --git a/examples/agent/startup-simulator-3000/agent_framework/utils/validation.py b/examples/agent/startup-simulator-3000/agent_framework/utils/validation.py index a23f9c86..7952e315 100644 --- a/examples/agent/startup-simulator-3000/agent_framework/utils/validation.py +++ b/examples/agent/startup-simulator-3000/agent_framework/utils/validation.py @@ -1,8 +1,7 @@ -import json -from datetime import datetime from typing import Any - +import json from agent_framework.llm.models import LLMMessage +from datetime import datetime def ensure_valid_io(data: Any) -> str: diff --git a/examples/agent/startup-simulator-3000/app.py b/examples/agent/startup-simulator-3000/app.py index a1c22bc6..9b403880 100644 --- a/examples/agent/startup-simulator-3000/app.py +++ b/examples/agent/startup-simulator-3000/app.py @@ -3,14 +3,13 @@ import asyncio import json import os - +from dotenv import load_dotenv +from flask import Flask, render_template, request, jsonify +from flask_cors import CORS from agent import SimpleAgent -from agent_framework.llm.models import LLMConfig from agent_framework.llm.openai_provider import OpenAIProvider +from agent_framework.llm.models import LLMConfig from agent_framework.utils.logging import get_galileo_logger -from dotenv import load_dotenv -from flask import Flask, jsonify, render_template, request -from flask_cors import CORS # Load environment variables load_dotenv() @@ -39,7 +38,12 @@ def generate_startup(): api_request = { "endpoint": "/api/generate", "method": "POST", - "inputs": {"industry": industry, "audience": audience, "random_word": random_word, "mode": mode}, + "inputs": { + "industry": industry, + "audience": audience, + "random_word": random_word, + "mode": mode, + }, } print(f"API Request: {json.dumps(api_request, indent=2)}") @@ -88,9 +92,7 @@ def generate_startup(): "status": "success", "result_length": len(final_output), "mode": mode, - "agent_result_preview": ( - str(agent_result)[:200] + "..." if len(str(agent_result)) > 200 else str(agent_result) - ), + "agent_result_preview": (str(agent_result)[:200] + "..." if len(str(agent_result)) > 200 else str(agent_result)), } print(f"API Response: {json.dumps(api_response, indent=2)}") @@ -161,7 +163,8 @@ async def run_agent(industry: str, audience: str, random_word: str, mode: str = ) # Run the agent with individual parameters (Splunk AO logging handled by individual traces) - return await agent.run(task, industry=industry, audience=audience, random_word=random_word) + result = await agent.run(task, industry=industry, audience=audience, random_word=random_word) + return result if __name__ == "__main__": diff --git a/examples/agent/startup-simulator-3000/demo.py b/examples/agent/startup-simulator-3000/demo.py index 66c6279f..30de2f94 100644 --- a/examples/agent/startup-simulator-3000/demo.py +++ b/examples/agent/startup-simulator-3000/demo.py @@ -7,17 +7,16 @@ import asyncio import json import os - +from dotenv import load_dotenv from agent import SimpleAgent -from agent_framework.llm.models import LLMConfig from agent_framework.llm.openai_provider import OpenAIProvider -from dotenv import load_dotenv +from agent_framework.llm.models import LLMConfig # Load environment variables load_dotenv() -async def demo_silly_mode() -> None: +async def demo_silly_mode(): """Demonstrate silly mode startup generation""" print("🎭 DEMO: Silly Mode Startup Generation") print("=" * 50) @@ -64,7 +63,7 @@ async def demo_silly_mode() -> None: print(f"❌ Error: {e}") -async def demo_serious_mode() -> None: +async def demo_serious_mode(): """Demonstrate serious mode startup generation""" print("\n💼 DEMO: Serious Mode Startup Generation") print("=" * 50) @@ -111,7 +110,7 @@ async def demo_serious_mode() -> None: print(f"❌ Error: {e}") -async def demo_individual_tools() -> None: +async def demo_individual_tools(): """Demonstrate individual tool usage""" print("\n🔧 DEMO: Individual Tool Usage") print("=" * 50) @@ -138,7 +137,7 @@ async def demo_individual_tools() -> None: print(f"❌ Error: {e}") -def check_environment() -> None: +def check_environment(): """Check if environment is properly configured""" print("🔍 Environment Check") print("=" * 50) @@ -166,7 +165,7 @@ def check_environment() -> None: print() -async def main() -> None: +async def main(): """Run all demos""" print("🚀 Startup Simulator 3000 - Demo Script") print("=" * 60) diff --git a/examples/agent/startup-simulator-3000/run_startup_sim.py b/examples/agent/startup-simulator-3000/run_startup_sim.py index 209c93d6..2c340650 100644 --- a/examples/agent/startup-simulator-3000/run_startup_sim.py +++ b/examples/agent/startup-simulator-3000/run_startup_sim.py @@ -1,18 +1,16 @@ import asyncio -import os - from agent import SimpleAgent -from agent_framework.llm.models import LLMConfig from agent_framework.llm.openai_provider import OpenAIProvider -from dotenv import load_dotenv - +from agent_framework.llm.models import LLMConfig from splunk_ao import splunk_ao_context +import os +from dotenv import load_dotenv # Load environment variables load_dotenv() -async def main() -> None: +async def main(): # Ensure Splunk AO environment variables are set if not os.getenv("SPLUNK_AO_API_KEY"): print("Warning: SPLUNK_AO_API_KEY not set. Splunk AO logging will be disabled.") diff --git a/examples/agent/startup-simulator-3000/tools/__init__.py b/examples/agent/startup-simulator-3000/tools/__init__.py index e1e636ce..02ec6709 100644 --- a/examples/agent/startup-simulator-3000/tools/__init__.py +++ b/examples/agent/startup-simulator-3000/tools/__init__.py @@ -1,7 +1,7 @@ """Tools used by the simple agent""" -from .hackernews_tool import HackerNewsTool -from .keyword_extraction import KeywordExtractorTool from .text_analysis import TextAnalyzerTool +from .keyword_extraction import KeywordExtractorTool +from .hackernews_tool import HackerNewsTool -__all__ = ["HackerNewsTool", "KeywordExtractorTool", "TextAnalyzerTool"] +__all__ = ["TextAnalyzerTool", "KeywordExtractorTool", "HackerNewsTool"] diff --git a/examples/agent/startup-simulator-3000/tools/hackernews_tool.py b/examples/agent/startup-simulator-3000/tools/hackernews_tool.py index 3ae5e1e3..9d4a6286 100644 --- a/examples/agent/startup-simulator-3000/tools/hackernews_tool.py +++ b/examples/agent/startup-simulator-3000/tools/hackernews_tool.py @@ -1,13 +1,13 @@ import json -from datetime import datetime - import aiohttp -from agent_framework.models import ToolMetadata +from splunk_ao import log # 🔍 Splunk AO decorator import for tool spans from agent_framework.tools.base import BaseTool -from agent_framework.utils.logging import get_galileo_logger # 🔍 Splunk AO helper import - gets centralized logger +from agent_framework.models import ToolMetadata +from agent_framework.utils.logging import ( + get_galileo_logger, +) # 🔍 Splunk AO helper import - gets centralized logger from dotenv import load_dotenv - -from splunk_ao import log # 🔍 Splunk AO decorator import for tool spans +from datetime import datetime # Load environment variables load_dotenv() @@ -32,10 +32,19 @@ def get_metadata(cls) -> ToolMetadata: tags=["news", "inspiration", "tech", "trending"], input_schema={ "type": "object", - "properties": {"limit": {"type": "integer", "description": "Number of stories to fetch", "default": 3}}, + "properties": { + "limit": { + "type": "integer", + "description": "Number of stories to fetch", + "default": 3, + } + }, "required": [], }, - output_schema={"type": "string", "description": "JSON string containing HackerNews stories with metadata"}, + output_schema={ + "type": "string", + "description": "JSON string containing HackerNews stories with metadata", + }, ) # 👀 GALILEO TOOL SPAN DECORATOR: This decorator creates a tool span for HTTP API calls @@ -79,9 +88,7 @@ async def execute(self, limit: int = 3) -> str: # Fetch individual story details stories = [] for story_id in story_ids: - async with session.get( - f"https://hacker-news.firebaseio.com/v0/item/{story_id}.json" - ) as story_response: + async with session.get(f"https://hacker-news.firebaseio.com/v0/item/{story_id}.json") as story_response: if story_response.status == 200: story_data = await story_response.json() if story_data and "title" in story_data: @@ -168,9 +175,7 @@ async def _execute_without_galileo(self, limit: int = 3) -> str: # Fetch individual story details stories = [] for story_id in story_ids: - async with session.get( - f"https://hacker-news.firebaseio.com/v0/item/{story_id}.json" - ) as story_response: + async with session.get(f"https://hacker-news.firebaseio.com/v0/item/{story_id}.json") as story_response: if story_response.status == 200: story_data = await story_response.json() if story_data and "title" in story_data: @@ -214,7 +219,7 @@ async def _execute_without_galileo(self, limit: int = 3) -> str: # ℹ️ TEST FUNCTION: This function can be used to test the tool independently -async def main() -> None: +async def main(): """Test the HackerNews tool""" tool = HackerNewsTool() result = await tool.execute(limit=3) diff --git a/examples/agent/startup-simulator-3000/tools/keyword_extraction.py b/examples/agent/startup-simulator-3000/tools/keyword_extraction.py index fb40da12..5d4bb6ea 100644 --- a/examples/agent/startup-simulator-3000/tools/keyword_extraction.py +++ b/examples/agent/startup-simulator-3000/tools/keyword_extraction.py @@ -1,7 +1,6 @@ -from typing import Any - -from agent_framework.models import ToolMetadata +from typing import Dict, Any from agent_framework.tools.base import BaseTool +from agent_framework.models import ToolMetadata class KeywordExtractorTool(BaseTool): @@ -16,19 +15,27 @@ def get_metadata(cls) -> ToolMetadata: tags=["text", "keywords", "extraction"], input_schema={ "type": "object", - "properties": {"text": {"type": "string", "description": "Text to extract keywords from"}}, + "properties": { + "text": { + "type": "string", + "description": "Text to extract keywords from", + } + }, "required": ["text"], }, output_schema={ "type": "object", "properties": { "keywords": {"type": "array", "items": {"type": "string"}}, - "importance_scores": {"type": "object", "additionalProperties": {"type": "number"}}, + "importance_scores": { + "type": "object", + "additionalProperties": {"type": "number"}, + }, }, }, ) - async def execute(self, text: str) -> dict[str, Any]: + async def execute(self, text: str) -> Dict[str, Any]: """Extract keywords from text""" # Simple implementation - in real world would use NLP words = text.lower().split() @@ -46,4 +53,7 @@ async def execute(self, text: str) -> dict[str, Any]: max_freq = max(freq for _, freq in keywords) if keywords else 1 importance_scores = {word: freq / max_freq for word, freq in keywords} - return {"keywords": [word for word, _ in keywords], "importance_scores": importance_scores} + return { + "keywords": [word for word, _ in keywords], + "importance_scores": importance_scores, + } diff --git a/examples/agent/startup-simulator-3000/tools/news_api_tool.py b/examples/agent/startup-simulator-3000/tools/news_api_tool.py index 55580480..9e493b44 100644 --- a/examples/agent/startup-simulator-3000/tools/news_api_tool.py +++ b/examples/agent/startup-simulator-3000/tools/news_api_tool.py @@ -1,14 +1,14 @@ -import json import os -from datetime import datetime - +import json import aiohttp -from agent_framework.models import ToolMetadata +from splunk_ao import log # 🔍 Splunk AO decorator import for tool spans from agent_framework.tools.base import BaseTool -from agent_framework.utils.logging import get_galileo_logger # 🔍 Splunk AO helper import - gets centralized logger +from agent_framework.models import ToolMetadata +from agent_framework.utils.logging import ( + get_galileo_logger, +) # 🔍 Splunk AO helper import - gets centralized logger from dotenv import load_dotenv - -from splunk_ao import log # 🔍 Splunk AO decorator import for tool spans +from datetime import datetime # Load environment variables load_dotenv() @@ -34,12 +34,23 @@ def get_metadata(cls) -> ToolMetadata: input_schema={ "type": "object", "properties": { - "category": {"type": "string", "description": "News category to fetch", "default": "business"}, - "limit": {"type": "integer", "description": "Number of articles to fetch", "default": 5}, + "category": { + "type": "string", + "description": "News category to fetch", + "default": "business", + }, + "limit": { + "type": "integer", + "description": "Number of articles to fetch", + "default": 5, + }, }, "required": [], }, - output_schema={"type": "string", "description": "JSON string containing news articles with metadata"}, + output_schema={ + "type": "string", + "description": "JSON string containing news articles with metadata", + }, ) # 👀 GALILEO TOOL SPAN DECORATOR: This decorator creates a tool span for HTTP API calls @@ -50,7 +61,11 @@ async def execute(self, category: str = "business", limit: int = 5) -> str: """Fetch business news from NewsAPI""" # Log inputs - inputs = {"category": category, "limit": limit, "timestamp": datetime.now().isoformat()} + inputs = { + "category": category, + "limit": limit, + "timestamp": datetime.now().isoformat(), + } print(f"News API Tool Inputs: {json.dumps(inputs, indent=2)}") # 👀 GALILEO LOGGER SETUP: Get the Splunk AO logger for this execution @@ -78,28 +93,34 @@ async def execute(self, category: str = "business", limit: int = 5) -> str: # Fetch news from NewsAPI url = "https://newsapi.org/v2/top-headlines" - params = {"country": "us", "category": category, "apiKey": api_key, "pageSize": limit} + params = { + "country": "us", + "category": category, + "apiKey": api_key, + "pageSize": limit, + } - async with aiohttp.ClientSession() as session, session.get(url, params=params) as response: - if response.status != 200: - raise Exception(f"Failed to fetch news: {response.status}") + async with aiohttp.ClientSession() as session: + async with session.get(url, params=params) as response: + if response.status != 200: + raise Exception(f"Failed to fetch news: {response.status}") - data = await response.json() + data = await response.json() - if data.get("status") != "ok": - raise Exception(f"NewsAPI error: {data.get('message', 'Unknown error')}") + if data.get("status") != "ok": + raise Exception(f"NewsAPI error: {data.get('message', 'Unknown error')}") - articles = data.get("articles", []) + articles = data.get("articles", []) - # Format articles for context - formatted_articles = [] - for article in articles[:limit]: - title = article.get("title", "No title") - description = article.get("description", "No description") - source = article.get("source", {}).get("name", "Unknown source") - formatted_articles.append(f"• {title} ({source}) - {description}") + # Format articles for context + formatted_articles = [] + for article in articles[:limit]: + title = article.get("title", "No title") + description = article.get("description", "No description") + source = article.get("source", {}).get("name", "Unknown source") + formatted_articles.append(f"• {title} ({source}) - {description}") - context = "\n".join(formatted_articles) + context = "\n".join(formatted_articles) # Create structured output output = { @@ -163,28 +184,34 @@ async def _execute_without_galileo(self, category: str = "business", limit: int # Fetch news from NewsAPI url = "https://newsapi.org/v2/top-headlines" - params = {"country": "us", "category": category, "apiKey": api_key, "pageSize": limit} + params = { + "country": "us", + "category": category, + "apiKey": api_key, + "pageSize": limit, + } - async with aiohttp.ClientSession() as session, session.get(url, params=params) as response: - if response.status != 200: - raise Exception(f"Failed to fetch news: {response.status}") + async with aiohttp.ClientSession() as session: + async with session.get(url, params=params) as response: + if response.status != 200: + raise Exception(f"Failed to fetch news: {response.status}") - data = await response.json() + data = await response.json() - if data.get("status") != "ok": - raise Exception(f"NewsAPI error: {data.get('message', 'Unknown error')}") + if data.get("status") != "ok": + raise Exception(f"NewsAPI error: {data.get('message', 'Unknown error')}") - articles = data.get("articles", []) + articles = data.get("articles", []) - # Format articles for context - formatted_articles = [] - for article in articles[:limit]: - title = article.get("title", "No title") - description = article.get("description", "No description") - source = article.get("source", {}).get("name", "Unknown source") - formatted_articles.append(f"• {title} ({source}) - {description}") + # Format articles for context + formatted_articles = [] + for article in articles[:limit]: + title = article.get("title", "No title") + description = article.get("description", "No description") + source = article.get("source", {}).get("name", "Unknown source") + formatted_articles.append(f"• {title} ({source}) - {description}") - context = "\n".join(formatted_articles) + context = "\n".join(formatted_articles) # Create structured output output = { @@ -209,7 +236,7 @@ async def _execute_without_galileo(self, category: str = "business", limit: int # ℹ️ TEST FUNCTION: This function can be used to test the tool independently -async def main() -> None: +async def main(): """Test the News API tool""" tool = NewsAPITool() result = await tool.execute(category="business", limit=3) diff --git a/examples/agent/startup-simulator-3000/tools/serious_startup_simulator.py b/examples/agent/startup-simulator-3000/tools/serious_startup_simulator.py index f857234e..d3fbc4ec 100644 --- a/examples/agent/startup-simulator-3000/tools/serious_startup_simulator.py +++ b/examples/agent/startup-simulator-3000/tools/serious_startup_simulator.py @@ -1,14 +1,17 @@ -import asyncio -import json import os -from datetime import datetime - -from agent_framework.models import ToolMetadata +import json +from splunk_ao.openai import ( + openai, +) # 🔍 Splunk AO-wrapped OpenAI client for automatic logging +from typing import Dict from agent_framework.tools.base import BaseTool -from agent_framework.utils.logging import get_galileo_logger # 🔍 Splunk AO helper import - gets centralized logger +from agent_framework.models import ToolMetadata +from agent_framework.utils.logging import ( + get_galileo_logger, +) # 🔍 Splunk AO helper import - gets centralized logger +import asyncio from dotenv import load_dotenv - -from splunk_ao.openai import openai # 🔍 Splunk AO-wrapped OpenAI client for automatic logging +from datetime import datetime # Load environment variables load_dotenv() @@ -24,9 +27,7 @@ class SeriousStartupSimulatorTool(BaseTool): def __init__(self): super().__init__() self.name = "serious_startup_simulator" - self.description = ( - "Generate a professional startup business plan based on industry, audience, and market trends" - ) + self.description = "Generate a professional startup business plan based on industry, audience, and market trends" # 👀 GALILEO INITIALIZATION: Get the centralized Splunk AO logger instance # This ensures all tools use the same Splunk AO configuration and connection self.galileo_logger = get_galileo_logger() @@ -40,10 +41,19 @@ def get_metadata(cls) -> ToolMetadata: input_schema={ "type": "object", "properties": { - "industry": {"type": "string", "description": "Industry for the startup"}, + "industry": { + "type": "string", + "description": "Industry for the startup", + }, "audience": {"type": "string", "description": "Target audience"}, - "random_word": {"type": "string", "description": "A word to incorporate"}, - "news_context": {"type": "string", "description": "Recent news context for market analysis"}, + "random_word": { + "type": "string", + "description": "A word to incorporate", + }, + "news_context": { + "type": "string", + "description": "Recent news context for market analysis", + }, }, "required": ["industry", "audience", "random_word"], }, @@ -128,9 +138,7 @@ async def execute(self, industry: str, audience: str, random_word: str, news_con "timestamp": datetime.now().isoformat(), "model": "gpt-4", "input_tokens": (response.usage.prompt_tokens if hasattr(response.usage, "prompt_tokens") else 0), - "output_tokens": ( - response.usage.completion_tokens if hasattr(response.usage, "completion_tokens") else 0 - ), + "output_tokens": (response.usage.completion_tokens if hasattr(response.usage, "completion_tokens") else 0), "total_tokens": (response.usage.total_tokens if hasattr(response.usage, "total_tokens") else 0), } @@ -187,9 +195,7 @@ async def execute(self, industry: str, audience: str, random_word: str, news_con raise e - async def _execute_without_galileo( - self, industry: str, audience: str, random_word: str, news_context: str = "" - ) -> str: + async def _execute_without_galileo(self, industry: str, audience: str, random_word: str, news_context: str = "") -> str: """Fallback execution without Splunk AO logging""" # ℹ️ FALLBACK METHOD: This method runs when Splunk AO is not available # It performs the same functionality but without any observability logging @@ -241,19 +247,25 @@ async def _execute_without_galileo( return json.dumps(galileo_output, indent=2) - def _parse_business_pitch(self, content: str) -> dict[str, str]: + def _parse_business_pitch(self, content: str) -> Dict[str, str]: """Parse the business pitch into structured components""" # This method could be enhanced to extract specific business plan sections # For now, it returns a simple structure - return {"executive_summary": (content[:200] + "..." if len(content) > 200 else content), "full_pitch": content} + return { + "executive_summary": (content[:200] + "..." if len(content) > 200 else content), + "full_pitch": content, + } # ℹ️ TEST FUNCTION: This function can be used to test the tool independently -async def main() -> None: +async def main(): """Test the Serious Startup Simulator tool""" tool = SeriousStartupSimulatorTool() result = await tool.execute( - industry="Finance", audience="Investors", random_word="fintech", news_context="Sample news context for testing" + industry="Finance", + audience="Investors", + random_word="fintech", + news_context="Sample news context for testing", ) print(f"Result: {result}") diff --git a/examples/agent/startup-simulator-3000/tools/startup_simulator.py b/examples/agent/startup-simulator-3000/tools/startup_simulator.py index a9baca88..2f7a8d49 100644 --- a/examples/agent/startup-simulator-3000/tools/startup_simulator.py +++ b/examples/agent/startup-simulator-3000/tools/startup_simulator.py @@ -1,14 +1,16 @@ -import asyncio -import json import os -from datetime import datetime - -from agent_framework.models import ToolMetadata +import json +from splunk_ao.openai import ( + openai, +) # 🔍 Splunk AO-wrapped OpenAI client for automatic logging from agent_framework.tools.base import BaseTool -from agent_framework.utils.logging import get_galileo_logger # 🔍 Splunk AO helper import - gets centralized logger +from agent_framework.models import ToolMetadata +from agent_framework.utils.logging import ( + get_galileo_logger, +) # 🔍 Splunk AO helper import - gets centralized logger +import asyncio from dotenv import load_dotenv - -from splunk_ao.openai import openai # 🔍 Splunk AO-wrapped OpenAI client for automatic logging +from datetime import datetime # Load environment variables load_dotenv() @@ -38,10 +40,19 @@ def get_metadata(cls) -> ToolMetadata: input_schema={ "type": "object", "properties": { - "industry": {"type": "string", "description": "Industry for the startup"}, + "industry": { + "type": "string", + "description": "Industry for the startup", + }, "audience": {"type": "string", "description": "Target audience"}, - "random_word": {"type": "string", "description": "A random word to include"}, - "hn_context": {"type": "string", "description": "Recent HackerNews context for inspiration"}, + "random_word": { + "type": "string", + "description": "A random word to include", + }, + "hn_context": { + "type": "string", + "description": "Recent HackerNews context for inspiration", + }, }, "required": ["industry", "audience", "random_word"], }, @@ -122,9 +133,7 @@ async def execute(self, industry: str, audience: str, random_word: str, hn_conte "timestamp": datetime.now().isoformat(), "model": "gpt-4", "input_tokens": (response.usage.prompt_tokens if hasattr(response.usage, "prompt_tokens") else 0), - "output_tokens": ( - response.usage.completion_tokens if hasattr(response.usage, "completion_tokens") else 0 - ), + "output_tokens": (response.usage.completion_tokens if hasattr(response.usage, "completion_tokens") else 0), "total_tokens": (response.usage.total_tokens if hasattr(response.usage, "total_tokens") else 0), } @@ -182,9 +191,7 @@ async def execute(self, industry: str, audience: str, random_word: str, hn_conte raise e - async def _execute_without_galileo( - self, industry: str, audience: str, random_word: str, hn_context: str = "" - ) -> str: + async def _execute_without_galileo(self, industry: str, audience: str, random_word: str, hn_context: str = "") -> str: """Fallback execution without Splunk AO logging""" # ℹ️ FALLBACK METHOD: This method runs when Splunk AO is not available # It performs the same functionality but without any observability logging @@ -234,11 +241,14 @@ async def _execute_without_galileo( # ℹ️ TEST FUNCTION: This function can be used to test the tool independently -async def main() -> None: +async def main(): """Test the Startup Simulator tool""" tool = StartupSimulatorTool() result = await tool.execute( - industry="Tech", audience="Developers", random_word="blockchain", hn_context="Sample HN context for testing" + industry="Tech", + audience="Developers", + random_word="blockchain", + hn_context="Sample HN context for testing", ) print(f"Result: {result}") diff --git a/examples/agent/startup-simulator-3000/tools/text_analysis.py b/examples/agent/startup-simulator-3000/tools/text_analysis.py index f2e51939..a6f40bb8 100644 --- a/examples/agent/startup-simulator-3000/tools/text_analysis.py +++ b/examples/agent/startup-simulator-3000/tools/text_analysis.py @@ -1,7 +1,6 @@ -from typing import Any - -from agent_framework.models import ToolMetadata +from typing import Dict, Any from agent_framework.tools.base import BaseTool +from agent_framework.models import ToolMetadata class TextAnalyzerTool(BaseTool): @@ -29,7 +28,7 @@ def get_metadata(cls) -> ToolMetadata: }, ) - async def execute(self, text: str) -> dict[str, Any]: + async def execute(self, text: str) -> Dict[str, Any]: """Analyze text complexity""" # Simple implementation - in real world would use NLP word_count = len(text.split()) diff --git a/examples/agent/strands-agents/agent.py b/examples/agent/strands-agents/agent.py index 42135a4b..c232f4b1 100644 --- a/examples/agent/strands-agents/agent.py +++ b/examples/agent/strands-agents/agent.py @@ -1,17 +1,16 @@ import os -# Load environment variables from the .env file -from dotenv import load_dotenv from strands import Agent, tool from strands.telemetry import StrandsTelemetry from strands_tools import calculator, current_time +# Load environment variables from the .env file +from dotenv import load_dotenv + load_dotenv(override=True) # Export the Splunk AO OTel API endpoint for OTel -os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = os.environ.get( - "SPLUNK_AO_API_ENDPOINT", "https://api.galileo.ai/otel/traces" -) +os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = os.environ.get("SPLUNK_AO_API_ENDPOINT", "https://api.galileo.ai/otel/traces") # Export the Splunk AO OTel headers pointing to the correct API key, project, and log stream headers = { diff --git a/examples/agent/weather-vibes-agent/__init__.py b/examples/agent/weather-vibes-agent/__init__.py index b791169c..0232f23f 100644 --- a/examples/agent/weather-vibes-agent/__init__.py +++ b/examples/agent/weather-vibes-agent/__init__.py @@ -10,11 +10,11 @@ # Import tools for easy access try: - from tools.recommendation_tool import RecommendationsTool from tools.weather_tool import WeatherTool + from tools.recommendation_tool import RecommendationsTool from tools.youtube_tool import YouTubeTool except ImportError: pass # Handle the case if the tools aren't ready yet # Define what's available from this package -__all__ = ["RecommendationsTool", "WeatherTool", "WeatherVibesAgent", "YouTubeTool"] +__all__ = ["WeatherVibesAgent", "WeatherTool", "RecommendationsTool", "YouTubeTool"] diff --git a/examples/agent/weather-vibes-agent/agent.py b/examples/agent/weather-vibes-agent/agent.py index 6a02a0cd..371263f2 100644 --- a/examples/agent/weather-vibes-agent/agent.py +++ b/examples/agent/weather-vibes-agent/agent.py @@ -8,14 +8,12 @@ python galileo_agent.py -l "Tokyo" -u imperial -m relaxing """ -import argparse import asyncio +import argparse import os import sys from pathlib import Path - from dotenv import load_dotenv - from splunk_ao import log, splunk_ao_context # Load environment variables & set up path @@ -23,7 +21,12 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # Quick environment check -required_keys = ["OPENAI_API_KEY", "WEATHERAPI_KEY", "YOUTUBE_API_KEY", "SPLUNK_AO_API_KEY"] +required_keys = [ + "OPENAI_API_KEY", + "WEATHERAPI_KEY", + "YOUTUBE_API_KEY", + "SPLUNK_AO_API_KEY", +] if any(not os.getenv(key) for key in required_keys): missing = [key for key in required_keys if not os.getenv(key)] print(f"Missing API keys: {', '.join(missing)}") @@ -45,19 +48,22 @@ @log(span_type="tool", name="weather_tool") async def get_weather(weather_tool, location, days=1): """Get weather data with Splunk AO tracing""" - return await weather_tool.execute(location=location, days=days) + result = await weather_tool.execute(location=location, days=days) + return result @log(span_type="tool", name="recommendations_tool") async def get_recommendations(recommendations_tool, weather, max_items=5): """Get recommendations with Splunk AO tracing""" - return await recommendations_tool.execute(weather=weather, max_items=max_items) + result = await recommendations_tool.execute(weather=weather, max_items=max_items) + return result @log(span_type="tool", name="youtube_tool") async def find_weather_video(youtube_tool, weather_condition, mood_override=None): """Find YouTube videos with Splunk AO tracing""" - return await youtube_tool.execute(weather_condition=weather_condition, mood_override=mood_override) + result = await youtube_tool.execute(weather_condition=weather_condition, mood_override=mood_override) + return result @log(span_type="workflow", name="weather_vibes_workflow") @@ -92,14 +98,21 @@ async def process_request(agent, request): # Execute tools weather_result = await get_weather(agent.weather_tool, location, days=1) if "error" in weather_result: - return {"error": 500, "message": f"Weather API error: {weather_result['message']}"} + return { + "error": 500, + "message": f"Weather API error: {weather_result['message']}", + } recommendations = await get_recommendations(agent.recommendations_tool, weather_result, max_recommendations) video_result = await find_weather_video(agent.youtube_tool, weather_result["condition"], video_mood) # Prepare response - result = {"weather": weather_result, "recommendations": recommendations, "video": video_result} + result = { + "weather": weather_result, + "recommendations": recommendations, + "video": video_result, + } # Filter weather details if not verbose if not verbose and "weather" in result: @@ -122,12 +135,12 @@ async def process_request(agent, request): return response except Exception as e: - return {"error": 500, "message": f"Error: {e!s}"} + return {"error": 500, "message": f"Error: {str(e)}"} # Log the inputs using the Splunk AO decorator @log(span_type="workflow", name="weather_vibes_agent") -async def run_agent_with_inputs(location, units, mood, recommendations, verbose) -> None: +async def run_agent_with_inputs(location, units, mood, recommendations, verbose): """Run the agent with specific inputs logged via the decorator""" print(f"Getting weather for: {location} (with Splunk AO tracing)") @@ -135,8 +148,16 @@ async def run_agent_with_inputs(location, units, mood, recommendations, verbose) agent = WeatherVibesAgent() request = { "input": {"location": location, "units": units}, - "config": {"verbose": verbose, "max_recommendations": recommendations, "video_mood": mood}, - "metadata": {"user_id": "demo_user", "session_id": "demo_session", "galileo_instrumented": True}, + "config": { + "verbose": verbose, + "max_recommendations": recommendations, + "video_mood": mood, + }, + "metadata": { + "user_id": "demo_user", + "session_id": "demo_session", + "galileo_instrumented": True, + }, } try: @@ -192,14 +213,23 @@ async def run_agent_with_inputs(location, units, mood, recommendations, verbose) # Use splunk_ao_context to set up the trace environment with further information -async def main() -> None: +async def main(): """Main entry point that uses splunk_ao_context to set up the trace environment""" # Parse arguments parser = argparse.ArgumentParser(description="Run the Splunk AO-instrumented Weather Vibes Agent") parser.add_argument("location", nargs="?", help="Location (e.g., 'Tokyo')") - parser.add_argument("-l", "--location", dest="location_alt", help="Alternative location specification") parser.add_argument( - "-u", "--units", choices=["metric", "imperial"], default="metric", help="Units (metric/imperial)" + "-l", + "--location", + dest="location_alt", + help="Alternative location specification", + ) + parser.add_argument( + "-u", + "--units", + choices=["metric", "imperial"], + default="metric", + help="Units (metric/imperial)", ) parser.add_argument("-m", "--mood", help="Video mood (e.g., 'relaxing', 'upbeat')") parser.add_argument("-r", "--recommendations", type=int, default=5, help="Number of recommendations") diff --git a/examples/agent/weather-vibes-agent/agent/descriptor.py b/examples/agent/weather-vibes-agent/agent/descriptor.py index f145e14a..e13cb125 100644 --- a/examples/agent/weather-vibes-agent/agent/descriptor.py +++ b/examples/agent/weather-vibes-agent/agent/descriptor.py @@ -13,7 +13,12 @@ "description": "An agent that provides weather information, item recommendations, and matching YouTube videos.", }, "specs": { - "capabilities": {"threads": True, "interrupts": False, "callbacks": False, "streaming": True}, + "capabilities": { + "threads": True, + "interrupts": False, + "callbacks": False, + "streaming": True, + }, "input": { "type": "object", "description": "Input for the Weather Vibes agent", @@ -76,7 +81,10 @@ "default": 5, "description": "Maximum number of recommendations to provide", }, - "video_mood": {"type": "string", "description": "Optional mood override for video selection"}, + "video_mood": { + "type": "string", + "description": "Optional mood override for video selection", + }, }, }, "thread_state": { diff --git a/examples/agent/weather-vibes-agent/agent/weather_vibes_agent.py b/examples/agent/weather-vibes-agent/agent/weather_vibes_agent.py index 66440f1b..72c53e56 100644 --- a/examples/agent/weather-vibes-agent/agent/weather_vibes_agent.py +++ b/examples/agent/weather-vibes-agent/agent/weather_vibes_agent.py @@ -2,17 +2,17 @@ Weather Vibes Agent implementation using the Simple Agent Framework. """ +import os import json import logging -import os import sys from pathlib import Path -from typing import Any +from typing import Dict, Any +from jinja2 import Environment, FileSystemLoader from agent_framework.agent import Agent -from agent_framework.models import ToolMetadata from agent_framework.state import AgentState -from jinja2 import Environment, FileSystemLoader +from agent_framework.models import ToolMetadata from openai import OpenAI # Configure proper path for imports @@ -21,10 +21,10 @@ sys.path.insert(0, str(project_root)) # import tools -from agent.descriptor import WEATHER_VIBES_DESCRIPTOR -from tools.recommendation_tool import RecommendationsTool from tools.weather_tool import WeatherTool +from tools.recommendation_tool import RecommendationsTool from tools.youtube_tool import YouTubeTool +from agent.descriptor import WEATHER_VIBES_DESCRIPTOR # Configure standard logging logger = logging.getLogger("weather_vibes_agent") @@ -48,7 +48,9 @@ def metadata(cls): # Add metadata method to the tool classes if they don't have it if not hasattr(WeatherTool, "metadata"): WeatherTool.metadata = create_tool_metadata( - "get_weather", "Get the current weather conditions for a location", ["weather", "utility"] + "get_weather", + "Get the current weather conditions for a location", + ["weather", "utility"], ) if not hasattr(RecommendationsTool, "metadata"): @@ -60,7 +62,9 @@ def metadata(cls): if not hasattr(YouTubeTool, "metadata"): YouTubeTool.metadata = create_tool_metadata( - "find_weather_video", "Find a YouTube video that matches the weather vibe", ["youtube", "entertainment"] + "find_weather_video", + "Find a YouTube video that matches the weather vibe", + ["youtube", "entertainment"], ) @@ -121,7 +125,8 @@ def _register_tools(self) -> None: self.tool_registry.register(metadata=WeatherTool.metadata(), implementation=self.weather_tool) self.tool_registry.register( - metadata=RecommendationsTool.metadata(), implementation=self.recommendations_tool + metadata=RecommendationsTool.metadata(), + implementation=self.recommendations_tool, ) self.tool_registry.register(metadata=YouTubeTool.metadata(), implementation=self.youtube_tool) @@ -140,7 +145,7 @@ async def _generate_system_prompt(self) -> str: favorite_locations=getattr(self.state, "favorite_locations", []), ) - async def _format_result(self, result: Any) -> dict[str, Any]: + async def _format_result(self, result: Any) -> Dict[str, Any]: """ Format the result from processing. @@ -155,23 +160,24 @@ async def _format_result(self, result: Any) -> dict[str, Any]: """ if isinstance(result, dict): return result - if hasattr(result, "model_dump"): + elif hasattr(result, "model_dump"): # Handle pydantic models return result.model_dump() - if hasattr(result, "__dict__"): + elif hasattr(result, "__dict__"): # Handle objects with __dict__ return result.__dict__ - # Default case - return {"result": str(result)} + else: + # Default case + return {"result": str(result)} - async def get_acp_descriptor(self) -> dict[str, Any]: + async def get_acp_descriptor(self) -> Dict[str, Any]: """ Return the ACP descriptor for this agent. This implements the ACP agent discovery capability. """ return self.descriptor - async def process_acp_request(self, request: dict[str, Any]) -> dict[str, Any]: + async def process_acp_request(self, request: Dict[str, Any]) -> Dict[str, Any]: """ Process an ACP request and generate a response. This implements the ACP run execution capability. @@ -201,7 +207,10 @@ async def process_acp_request(self, request: dict[str, Any]) -> dict[str, Any]: # Validate input if not location: logger.error("Invalid input: 'location' field is required") - return {"error": 400, "message": "Invalid input: 'location' field is required"} + return { + "error": 400, + "message": "Invalid input: 'location' field is required", + } # Update search history if not hasattr(self.state, "search_history"): @@ -219,22 +228,25 @@ async def process_acp_request(self, request: dict[str, Any]) -> dict[str, Any]: weather_result = await self.weather_tool.execute(location=location, days=1) if "error" in weather_result: logger.error(f"Weather API error: {weather_result['message']}") - return {"error": 500, "message": f"Weather API error: {weather_result['message']}"} + return { + "error": 500, + "message": f"Weather API error: {weather_result['message']}", + } # Step 2: Get recommendations logger.info("Getting recommendations based on weather") - recommendations = await self.recommendations_tool.execute( - weather=weather_result, max_items=max_recommendations - ) + recommendations = await self.recommendations_tool.execute(weather=weather_result, max_items=max_recommendations) # Step 3: Get matching YouTube video logger.info(f"Finding YouTube video matching weather condition: {weather_result['condition']}") - video_result = await self.youtube_tool.execute( - weather_condition=weather_result["condition"], mood_override=video_mood - ) + video_result = await self.youtube_tool.execute(weather_condition=weather_result["condition"], mood_override=video_mood) # Prepare the response - result = {"weather": weather_result, "recommendations": recommendations, "video": video_result} + result = { + "weather": weather_result, + "recommendations": recommendations, + "video": video_result, + } # If not verbose, filter out some weather details if not verbose and "weather" in result: @@ -261,5 +273,5 @@ async def process_acp_request(self, request: dict[str, Any]) -> dict[str, Any]: return response except Exception as e: - logger.error(f"Error processing request: {e!s}") - return {"error": 500, "message": f"Error processing request: {e!s}"} + logger.error(f"Error processing request: {str(e)}") + return {"error": 500, "message": f"Error processing request: {str(e)}"} diff --git a/examples/agent/weather-vibes-agent/descriptor.py b/examples/agent/weather-vibes-agent/descriptor.py index c8d2c201..d4bb37e0 100644 --- a/examples/agent/weather-vibes-agent/descriptor.py +++ b/examples/agent/weather-vibes-agent/descriptor.py @@ -13,7 +13,12 @@ "description": "An agent that provides weather information, item recommendations, and matching YouTube videos.", }, "specs": { - "capabilities": {"threads": True, "interrupts": False, "callbacks": False, "streaming": True}, + "capabilities": { + "threads": True, + "interrupts": False, + "callbacks": False, + "streaming": True, + }, "input": { "type": "object", "description": "Input for the Weather Vibes agent", @@ -76,7 +81,10 @@ "default": 5, "description": "Maximum number of recommendations to provide", }, - "video_mood": {"type": "string", "description": "Optional mood override for video selection"}, + "video_mood": { + "type": "string", + "description": "Optional mood override for video selection", + }, }, }, "thread_state": { diff --git a/examples/agent/weather-vibes-agent/tools/__init__.py b/examples/agent/weather-vibes-agent/tools/__init__.py index 02d71022..8fab87da 100644 --- a/examples/agent/weather-vibes-agent/tools/__init__.py +++ b/examples/agent/weather-vibes-agent/tools/__init__.py @@ -1,5 +1,5 @@ -from .recommendation_tool import RecommendationsTool from .weather_tool import WeatherTool +from .recommendation_tool import RecommendationsTool from .youtube_tool import YouTubeTool -__all__ = ["RecommendationsTool", "WeatherTool", "YouTubeTool"] +__all__ = ["WeatherTool", "RecommendationsTool", "YouTubeTool"] diff --git a/examples/agent/weather-vibes-agent/tools/recommendation_tool.py b/examples/agent/weather-vibes-agent/tools/recommendation_tool.py index 2fd4f5db..7246353e 100644 --- a/examples/agent/weather-vibes-agent/tools/recommendation_tool.py +++ b/examples/agent/weather-vibes-agent/tools/recommendation_tool.py @@ -2,16 +2,15 @@ Tool for generating item recommendations based on weather conditions. """ -from typing import Any - -from agent_framework.tools.base import BaseTool +from typing import Dict, Any, List from pydantic import BaseModel +from agent_framework.tools.base import BaseTool class RecommendationsInput(BaseModel): """Input schema for recommendations tool""" - weather: dict[str, Any] + weather: Dict[str, Any] max_items: int = 5 @@ -23,7 +22,7 @@ class RecommendationsTool(BaseTool): tags = ["weather", "recommendations"] input_schema = RecommendationsInput.model_json_schema() - async def execute(self, weather: dict[str, Any], max_items: int = 5) -> list[str]: + async def execute(self, weather: Dict[str, Any], max_items: int = 5) -> List[str]: """ Execute the tool to get recommendations. diff --git a/examples/agent/weather-vibes-agent/tools/weather_tool.py b/examples/agent/weather-vibes-agent/tools/weather_tool.py index 1f241c76..8dd9de0c 100644 --- a/examples/agent/weather-vibes-agent/tools/weather_tool.py +++ b/examples/agent/weather-vibes-agent/tools/weather_tool.py @@ -3,11 +3,10 @@ """ import os -from typing import Any - -import requests -from agent_framework.tools.base import BaseTool +from typing import Dict, Any from pydantic import BaseModel +from agent_framework.tools.base import BaseTool +import requests class WeatherInput(BaseModel): @@ -31,7 +30,7 @@ def __init__(self): raise ValueError("WeatherAPI.com API key not found in environment") self.base_url = "http://api.weatherapi.com/v1/forecast.json" - async def execute(self, location: str, days: int = 1) -> dict[str, Any]: + async def execute(self, location: str, days: int = 1) -> Dict[str, Any]: """ Execute the tool to get current weather and forecast. @@ -97,4 +96,7 @@ async def execute(self, location: str, days: int = 1) -> dict[str, Any]: return weather_info except Exception as e: - return {"error": str(e), "message": f"Failed to get weather for location: {location}"} + return { + "error": str(e), + "message": f"Failed to get weather for location: {location}", + } diff --git a/examples/agent/weather-vibes-agent/tools/youtube_tool.py b/examples/agent/weather-vibes-agent/tools/youtube_tool.py index 646e7a26..4c16dfce 100644 --- a/examples/agent/weather-vibes-agent/tools/youtube_tool.py +++ b/examples/agent/weather-vibes-agent/tools/youtube_tool.py @@ -3,18 +3,17 @@ """ import os -from typing import Any - +from typing import Dict, Any, Optional +from pydantic import BaseModel from agent_framework.tools.base import BaseTool from googleapiclient.discovery import build -from pydantic import BaseModel class YouTubeInput(BaseModel): """Input schema for YouTube tool""" weather_condition: str - mood_override: str | None = None + mood_override: Optional[str] = None class YouTubeTool(BaseTool): @@ -31,7 +30,7 @@ def __init__(self): raise ValueError("YouTube API key not found in environment") self.youtube = build("youtube", "v3", developerKey=self.api_key) - async def execute(self, weather_condition: str, mood_override: str | None = None) -> dict[str, Any]: + async def execute(self, weather_condition: str, mood_override: Optional[str] = None) -> Dict[str, Any]: """ Execute the tool to find a weather-matching YouTube video. @@ -83,7 +82,11 @@ async def execute(self, weather_condition: str, mood_override: str | None = None "channel": video["snippet"]["channelTitle"], "query": query, } - return {"error": "No videos found", "query": query} + else: + return {"error": "No videos found", "query": query} except Exception as e: - return {"error": str(e), "message": "Failed to find a matching YouTube video"} + return { + "error": str(e), + "message": "Failed to find a matching YouTube video", + } diff --git a/examples/chatbot/basic-examples/app.py b/examples/chatbot/basic-examples/app.py index a61ab8ca..35f8b5a8 100644 --- a/examples/chatbot/basic-examples/app.py +++ b/examples/chatbot/basic-examples/app.py @@ -1,14 +1,19 @@ import os +from splunk_ao import openai # The Splunk AO OpenAI client wrapper is all you need! from dotenv import load_dotenv -from splunk_ao import openai # The Splunk AO OpenAI client wrapper is all you need! - load_dotenv() -client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"), organization=os.environ.get("OPENAI_ORGANIZATION")) +client = openai.OpenAI( + api_key=os.environ.get("OPENAI_API_KEY"), + organization=os.environ.get("OPENAI_ORGANIZATION"), +) prompt = "Explain the following topic succinctly: Newton's First Law" -response = client.chat.completions.create(model="gpt-4o", messages=[{"role": "system", "content": prompt}]) +response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "system", "content": prompt}], +) print(response.choices[0].message.content.strip()) diff --git a/examples/chatbot/basic-examples/context.py b/examples/chatbot/basic-examples/context.py index 0e5b412a..b7a168fb 100644 --- a/examples/chatbot/basic-examples/context.py +++ b/examples/chatbot/basic-examples/context.py @@ -1,6 +1,6 @@ import os -from splunk_ao import openai, splunk_ao_context +from splunk_ao import splunk_ao_context, openai # If you've set your SPLUNK_AO_PROJECT and SPLUNK_AO_LOG_STREAM env vars, you can skip this step splunk_ao_context.init(project="your-project-id", log_stream="your-log-stream-id") @@ -10,9 +10,7 @@ def call_openai(): - chat_completion = client.chat.completions.create( - messages=[{"role": "user", "content": "Say this is a test"}], model="gpt-4o" - ) + chat_completion = client.chat.completions.create(messages=[{"role": "user", "content": "Say this is a test"}], model="gpt-4o") return chat_completion.choices[0].message.content diff --git a/examples/chatbot/basic-examples/hallucination.py b/examples/chatbot/basic-examples/hallucination.py index 96bdb0bf..053d98f8 100644 --- a/examples/chatbot/basic-examples/hallucination.py +++ b/examples/chatbot/basic-examples/hallucination.py @@ -1,11 +1,9 @@ -import json import os +import json import time - import openai # Using the standard OpenAI library -from dotenv import load_dotenv - from splunk_ao import SplunkAOLogger # Import SplunkAOLogger for logging +from dotenv import load_dotenv load_dotenv() @@ -13,7 +11,10 @@ logger = SplunkAOLogger(project="hallucination", log_stream="dev") # Initialize the standard OpenAI client -client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"), organization=os.environ.get("OPENAI_ORGANIZATION")) +client = openai.OpenAI( + api_key=os.environ.get("OPENAI_API_KEY"), + organization=os.environ.get("OPENAI_ORGANIZATION"), +) # Questions that might lead to hallucinations QUESTIONS = [ @@ -32,7 +33,11 @@ def generate_response(question: str, system_prompt: str) -> str: start_time = time.time_ns() response = client.chat.completions.create( - model="gpt-4o", messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": question}] + model="gpt-4o", + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": question}, + ], ) # Record the end time for the LLM call @@ -64,7 +69,7 @@ def run_hallucination_demo(): question_results = {"question": question, "responses": []} - print(f"\n\n{'=' * 80}\nQUESTION: {question}\n{'=' * 80}") + print(f"\n\n{'='*80}\nQUESTION: {question}\n{'='*80}") for prompt_name, system_prompt in SYSTEM_PROMPTS.items(): # Generate response diff --git a/examples/chatbot/basic-examples/test.py b/examples/chatbot/basic-examples/test.py index fddd000a..2b03d9d8 100644 --- a/examples/chatbot/basic-examples/test.py +++ b/examples/chatbot/basic-examples/test.py @@ -1,18 +1,20 @@ import os import time - import openai # Using the standard OpenAI library -from dotenv import load_dotenv - from splunk_ao import SplunkAOLogger # Import SplunkAOLogger for logging +from dotenv import load_dotenv + load_dotenv() # Initialize the SplunkAOLogger logger = SplunkAOLogger(project="chatbot", log_stream="test") # Initialize the standard OpenAI client -client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"), organization=os.environ.get("OPENAI_ORGANIZATION")) +client = openai.OpenAI( + api_key=os.environ.get("OPENAI_API_KEY"), + organization=os.environ.get("OPENAI_ORGANIZATION"), +) # Start a trace # prompt = f"Explain the following topic succinctly: Newton's First Law" @@ -27,7 +29,10 @@ start_time = time.time_ns() # Make the OpenAI API call directly -response = client.chat.completions.create(model="gpt-4o", messages=[{"role": "system", "content": prompt}]) +response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "system", "content": prompt}], +) # Record the end time for the LLM call end_time = time.time_ns() diff --git a/examples/chatbot/elevenlabs-chatbot/main.py b/examples/chatbot/elevenlabs-chatbot/main.py index 7b73da9a..d83a84c3 100644 --- a/examples/chatbot/elevenlabs-chatbot/main.py +++ b/examples/chatbot/elevenlabs-chatbot/main.py @@ -70,7 +70,7 @@ def on_user_transcript(transcript: str) -> None: # ============================================================================= -def run_voice_conversation() -> None: +def run_voice_conversation(): """Run a voice conversation with ElevenLabs + Splunk AO logging. This function: diff --git a/examples/chatbot/elevenlabs-chatbot/splunk_ao_handler.py b/examples/chatbot/elevenlabs-chatbot/splunk_ao_handler.py index 00cd0d00..a50b4215 100644 --- a/examples/chatbot/elevenlabs-chatbot/splunk_ao_handler.py +++ b/examples/chatbot/elevenlabs-chatbot/splunk_ao_handler.py @@ -10,8 +10,9 @@ """ import os +from typing import Optional -from splunk_ao import Message, MessageRole, SplunkAOLogger +from splunk_ao import SplunkAOLogger, Message, MessageRole class SplunkAOHandler: @@ -22,8 +23,8 @@ class SplunkAOHandler: """ def __init__(self): - self._logger: SplunkAOLogger | None = None - self._session_id: str | None = None + self._logger: Optional[SplunkAOLogger] = None + self._session_id: Optional[str] = None self._turn_count = 0 # Load Splunk AO config from environment @@ -32,7 +33,7 @@ def __init__(self): self._init_logger() - def _init_logger(self) -> None: + def _init_logger(self): """Initialize the Splunk AO Logger. The logger connects to a specific project and log stream in Splunk AO. @@ -40,12 +41,15 @@ def _init_logger(self) -> None: or by feature area. """ try: - self._logger = SplunkAOLogger(project=self._project_name, log_stream=self._log_stream) + self._logger = SplunkAOLogger( + project=self._project_name, + log_stream=self._log_stream, + ) print(f"[SPLUNK_AO] Logger initialized for project: {self._project_name}") except Exception as e: print(f"[SPLUNK_AO] Logger init failed: {e}") - def start_conversation(self, session_id: str) -> None: + def start_conversation(self, session_id: str): """Start a new conversation session in Splunk AO. A session groups all the turns of a single conversation together, @@ -102,7 +106,7 @@ def log_agent_turn(self, response: str) -> None: except Exception as e: print(f"[SPLUNK_AO] Logging error: {e}") - def end_conversation(self) -> None: + def end_conversation(self): """End the conversation session and cleanup. Ensures all logs are flushed and the session is properly closed. @@ -120,7 +124,7 @@ def end_conversation(self) -> None: # Singleton instance for the Splunk AO handler -_splunk_ao_handler: SplunkAOHandler | None = None +_splunk_ao_handler: Optional[SplunkAOHandler] = None def get_splunk_ao_handler() -> SplunkAOHandler: diff --git a/examples/chatbot/sample-project-chatbot/anthropic/app.py b/examples/chatbot/sample-project-chatbot/anthropic/app.py index b3b4983a..0b5b7128 100644 --- a/examples/chatbot/sample-project-chatbot/anthropic/app.py +++ b/examples/chatbot/sample-project-chatbot/anthropic/app.py @@ -25,13 +25,14 @@ """ -import os from datetime import datetime +import os from anthropic import Anthropic + from dotenv import load_dotenv -from splunk_ao import log, splunk_ao_context +from splunk_ao import splunk_ao_context, log # Load the environment variables from the .env file # This will override any existing environment variables with the same name @@ -101,7 +102,10 @@ def send_chat_to_anthropic() -> str: # Send the chat history to the Anthropic API and get the response response = client.messages.create( - max_tokens=1024, messages=chat_history_anthropic, system=system_prompt, model=MODEL_NAME + max_tokens=1024, + messages=chat_history_anthropic, + system=system_prompt, + model=MODEL_NAME, ) # Print the response to the console diff --git a/examples/chatbot/sample-project-chatbot/anthropic/test.py b/examples/chatbot/sample-project-chatbot/anthropic/test.py index 9208f1f0..50412891 100644 --- a/examples/chatbot/sample-project-chatbot/anthropic/test.py +++ b/examples/chatbot/sample-project-chatbot/anthropic/test.py @@ -10,15 +10,16 @@ import os import time -from app import chat_with_llm from dotenv import load_dotenv from splunk_ao import SplunkAOMetrics from splunk_ao.datasets import create_dataset, get_dataset from splunk_ao.experiments import get_experiment, run_experiment +from app import chat_with_llm + -def setup_module() -> None: +def setup_module(): """ Setup function that runs once when this test module is executed. @@ -35,7 +36,9 @@ def setup_module() -> None: raise ValueError("SPLUNK_AO_PROJECT and SPLUNK_AO_API_KEY environment variables are required") # Check to see if we already have the dataset, if not we can create it. - dataset = get_dataset(name="simple-chatbot-unit-test-dataset") + dataset = get_dataset( + name="simple-chatbot-unit-test-dataset", + ) # If we don't have the dataset, create it with some canned data. Some of these questions # are designed to be factual, while others are designed to be nonsensical or not @@ -43,13 +46,16 @@ def setup_module() -> None: # of the model when running the experiment. if dataset is None: # Load dataset content from JSON file - with open("../dataset.json", encoding="utf-8") as f: + with open("../dataset.json", "r", encoding="utf-8") as f: dataset_content = json.load(f) - dataset = create_dataset(name="simple-chatbot-unit-test-dataset", content=dataset_content) + dataset = create_dataset( + name="simple-chatbot-unit-test-dataset", + content=dataset_content, + ) -def test_run_experiment_with_dataset() -> None: +def test_run_experiment_with_dataset(): """ This test will run the dataset against our chatbot app using the Splunk AO experiments framework. This is designed to show how you can run experiments with a dataset against a real-world application @@ -68,7 +74,10 @@ def test_run_experiment_with_dataset() -> None: experiment_name="simple-chatbot-experiment", dataset_name="simple-chatbot-unit-test-dataset", function=chat_with_llm, - metrics=[SplunkAOMetrics.correctness, SplunkAOMetrics.instruction_adherence], + metrics=[ + SplunkAOMetrics.correctness, + SplunkAOMetrics.instruction_adherence, + ], project=os.getenv("SPLUNK_AO_PROJECT"), ) @@ -76,7 +85,8 @@ def test_run_experiment_with_dataset() -> None: # We need to use the project ID and experiment name from the response as we only have the project name, and the experiment name # is generated with a timestamp, so we can't use the name directly. experiment = get_experiment( - project_id=experiment_response["experiment"].project_id, experiment_name=experiment_response["experiment"].name + project_id=experiment_response["experiment"].project_id, + experiment_name=experiment_response["experiment"].name, ) while ( experiment.aggregate_metrics is None diff --git a/examples/chatbot/sample-project-chatbot/azure-inference/app.py b/examples/chatbot/sample-project-chatbot/azure-inference/app.py index 3f6d428f..dcf9f903 100644 --- a/examples/chatbot/sample-project-chatbot/azure-inference/app.py +++ b/examples/chatbot/sample-project-chatbot/azure-inference/app.py @@ -26,15 +26,16 @@ """ -import os from datetime import datetime +import os from azure.ai.inference import ChatCompletionsClient -from azure.ai.inference.models import AssistantMessage, SystemMessage, UserMessage +from azure.ai.inference.models import SystemMessage, UserMessage, AssistantMessage from azure.core.credentials import AzureKeyCredential + from dotenv import load_dotenv -from splunk_ao import log, splunk_ao_context +from splunk_ao import splunk_ao_context, log # Load the environment variables from the .env file # This will override any existing environment variables with the same name diff --git a/examples/chatbot/sample-project-chatbot/azure-inference/test.py b/examples/chatbot/sample-project-chatbot/azure-inference/test.py index 9208f1f0..50412891 100644 --- a/examples/chatbot/sample-project-chatbot/azure-inference/test.py +++ b/examples/chatbot/sample-project-chatbot/azure-inference/test.py @@ -10,15 +10,16 @@ import os import time -from app import chat_with_llm from dotenv import load_dotenv from splunk_ao import SplunkAOMetrics from splunk_ao.datasets import create_dataset, get_dataset from splunk_ao.experiments import get_experiment, run_experiment +from app import chat_with_llm + -def setup_module() -> None: +def setup_module(): """ Setup function that runs once when this test module is executed. @@ -35,7 +36,9 @@ def setup_module() -> None: raise ValueError("SPLUNK_AO_PROJECT and SPLUNK_AO_API_KEY environment variables are required") # Check to see if we already have the dataset, if not we can create it. - dataset = get_dataset(name="simple-chatbot-unit-test-dataset") + dataset = get_dataset( + name="simple-chatbot-unit-test-dataset", + ) # If we don't have the dataset, create it with some canned data. Some of these questions # are designed to be factual, while others are designed to be nonsensical or not @@ -43,13 +46,16 @@ def setup_module() -> None: # of the model when running the experiment. if dataset is None: # Load dataset content from JSON file - with open("../dataset.json", encoding="utf-8") as f: + with open("../dataset.json", "r", encoding="utf-8") as f: dataset_content = json.load(f) - dataset = create_dataset(name="simple-chatbot-unit-test-dataset", content=dataset_content) + dataset = create_dataset( + name="simple-chatbot-unit-test-dataset", + content=dataset_content, + ) -def test_run_experiment_with_dataset() -> None: +def test_run_experiment_with_dataset(): """ This test will run the dataset against our chatbot app using the Splunk AO experiments framework. This is designed to show how you can run experiments with a dataset against a real-world application @@ -68,7 +74,10 @@ def test_run_experiment_with_dataset() -> None: experiment_name="simple-chatbot-experiment", dataset_name="simple-chatbot-unit-test-dataset", function=chat_with_llm, - metrics=[SplunkAOMetrics.correctness, SplunkAOMetrics.instruction_adherence], + metrics=[ + SplunkAOMetrics.correctness, + SplunkAOMetrics.instruction_adherence, + ], project=os.getenv("SPLUNK_AO_PROJECT"), ) @@ -76,7 +85,8 @@ def test_run_experiment_with_dataset() -> None: # We need to use the project ID and experiment name from the response as we only have the project name, and the experiment name # is generated with a timestamp, so we can't use the name directly. experiment = get_experiment( - project_id=experiment_response["experiment"].project_id, experiment_name=experiment_response["experiment"].name + project_id=experiment_response["experiment"].project_id, + experiment_name=experiment_response["experiment"].name, ) while ( experiment.aggregate_metrics is None diff --git a/examples/chatbot/sample-project-chatbot/openai-ollama/app.py b/examples/chatbot/sample-project-chatbot/openai-ollama/app.py index a463efe1..9b6fa658 100644 --- a/examples/chatbot/sample-project-chatbot/openai-ollama/app.py +++ b/examples/chatbot/sample-project-chatbot/openai-ollama/app.py @@ -28,12 +28,12 @@ """ -import os from datetime import datetime +import os from dotenv import load_dotenv -from splunk_ao import log, splunk_ao_context +from splunk_ao import splunk_ao_context, log from splunk_ao.openai import OpenAI # Load the environment variables from the .env file diff --git a/examples/chatbot/sample-project-chatbot/openai-ollama/create_sample_logs.py b/examples/chatbot/sample-project-chatbot/openai-ollama/create_sample_logs.py index 3696f187..384afe8c 100644 --- a/examples/chatbot/sample-project-chatbot/openai-ollama/create_sample_logs.py +++ b/examples/chatbot/sample-project-chatbot/openai-ollama/create_sample_logs.py @@ -1,20 +1,20 @@ # A script to generate log streams # Load the dataset.json +from datetime import datetime import json import uuid -from datetime import datetime -from app import chat_history, chat_with_llm +from splunk_ao import splunk_ao_context + +from app import chat_with_llm, chat_history # Load environment variables from .env file from dotenv import load_dotenv -from splunk_ao import splunk_ao_context - load_dotenv(override=True) -with open("../dataset.json", encoding="utf-8") as f: +with open("../dataset.json", "r", encoding="utf-8") as f: dataset_content = json.load(f) print(f"Starting to log {len(dataset_content)} interactions...") diff --git a/examples/chatbot/sample-project-chatbot/openai-ollama/test.py b/examples/chatbot/sample-project-chatbot/openai-ollama/test.py index 9208f1f0..50412891 100644 --- a/examples/chatbot/sample-project-chatbot/openai-ollama/test.py +++ b/examples/chatbot/sample-project-chatbot/openai-ollama/test.py @@ -10,15 +10,16 @@ import os import time -from app import chat_with_llm from dotenv import load_dotenv from splunk_ao import SplunkAOMetrics from splunk_ao.datasets import create_dataset, get_dataset from splunk_ao.experiments import get_experiment, run_experiment +from app import chat_with_llm + -def setup_module() -> None: +def setup_module(): """ Setup function that runs once when this test module is executed. @@ -35,7 +36,9 @@ def setup_module() -> None: raise ValueError("SPLUNK_AO_PROJECT and SPLUNK_AO_API_KEY environment variables are required") # Check to see if we already have the dataset, if not we can create it. - dataset = get_dataset(name="simple-chatbot-unit-test-dataset") + dataset = get_dataset( + name="simple-chatbot-unit-test-dataset", + ) # If we don't have the dataset, create it with some canned data. Some of these questions # are designed to be factual, while others are designed to be nonsensical or not @@ -43,13 +46,16 @@ def setup_module() -> None: # of the model when running the experiment. if dataset is None: # Load dataset content from JSON file - with open("../dataset.json", encoding="utf-8") as f: + with open("../dataset.json", "r", encoding="utf-8") as f: dataset_content = json.load(f) - dataset = create_dataset(name="simple-chatbot-unit-test-dataset", content=dataset_content) + dataset = create_dataset( + name="simple-chatbot-unit-test-dataset", + content=dataset_content, + ) -def test_run_experiment_with_dataset() -> None: +def test_run_experiment_with_dataset(): """ This test will run the dataset against our chatbot app using the Splunk AO experiments framework. This is designed to show how you can run experiments with a dataset against a real-world application @@ -68,7 +74,10 @@ def test_run_experiment_with_dataset() -> None: experiment_name="simple-chatbot-experiment", dataset_name="simple-chatbot-unit-test-dataset", function=chat_with_llm, - metrics=[SplunkAOMetrics.correctness, SplunkAOMetrics.instruction_adherence], + metrics=[ + SplunkAOMetrics.correctness, + SplunkAOMetrics.instruction_adherence, + ], project=os.getenv("SPLUNK_AO_PROJECT"), ) @@ -76,7 +85,8 @@ def test_run_experiment_with_dataset() -> None: # We need to use the project ID and experiment name from the response as we only have the project name, and the experiment name # is generated with a timestamp, so we can't use the name directly. experiment = get_experiment( - project_id=experiment_response["experiment"].project_id, experiment_name=experiment_response["experiment"].name + project_id=experiment_response["experiment"].project_id, + experiment_name=experiment_response["experiment"].name, ) while ( experiment.aggregate_metrics is None diff --git a/examples/dataset-experiments/app-simple.py b/examples/dataset-experiments/app-simple.py index d41d2930..943c4805 100644 --- a/examples/dataset-experiments/app-simple.py +++ b/examples/dataset-experiments/app-simple.py @@ -1,9 +1,8 @@ -from dotenv import load_dotenv - from splunk_ao import Message, MessageRole -from splunk_ao.datasets import get_dataset +from splunk_ao.prompts import get_prompt, create_prompt from splunk_ao.experiments import run_experiment -from splunk_ao.prompts import create_prompt, get_prompt +from splunk_ao.datasets import get_dataset +from dotenv import load_dotenv load_dotenv() diff --git a/examples/dataset-experiments/app.py b/examples/dataset-experiments/app.py index 4606101d..0d47331f 100644 --- a/examples/dataset-experiments/app.py +++ b/examples/dataset-experiments/app.py @@ -7,18 +7,16 @@ import argparse from datetime import datetime - -from dotenv import load_dotenv - from splunk_ao import Message, MessageRole -from splunk_ao.datasets import get_dataset +from splunk_ao.prompts import get_prompt, create_prompt from splunk_ao.experiments import run_experiment -from splunk_ao.prompts import create_prompt, get_prompt +from splunk_ao.datasets import get_dataset +from dotenv import load_dotenv load_dotenv() -def main() -> None: +def main(): """ Main function to run Splunk AO experiments with command-line parameters. @@ -36,10 +34,16 @@ def main() -> None: ) parser.add_argument("--project", default="test project", help="Project name") parser.add_argument("--experiment", default="test experiment", help="Experiment name") - parser.add_argument("--prompt-name", default="default", help='Prompt name (optional, defaults to "default")') + parser.add_argument( + "--prompt-name", + default="default", + help='Prompt name (optional, defaults to "default")', + ) parser.add_argument("--dataset", default="default", help="Dataset name") parser.add_argument( - "--prompt-content", default="You are an assistant. Respond to the user input.", help="Prompt content" + "--prompt-content", + default="You are an assistant. Respond to the user input.", + help="Prompt content", ) args = parser.parse_args() project = args.project @@ -53,7 +57,10 @@ def main() -> None: prompt = create_prompt( name=prompt_name, template=[ - Message(role=MessageRole.system, content=prompt_content), + Message( + role=MessageRole.system, + content=prompt_content, + ), Message(role=MessageRole.user, content="{{input}}"), ], ) @@ -67,7 +74,10 @@ def main() -> None: prompt = create_prompt( name=new_prompt_name, template=[ - Message(role=MessageRole.system, content=prompt_content), + Message( + role=MessageRole.system, + content=prompt_content, + ), Message(role=MessageRole.user, content="{{input}}"), ], ) @@ -80,7 +90,11 @@ def main() -> None: experiment_name, dataset=dataset, prompt_template=prompt, - prompt_settings={"max_tokens": 256, "model_alias": "GPT-4o", "temperature": 0.8}, + prompt_settings={ + "max_tokens": 256, + "model_alias": "GPT-4o", + "temperature": 0.8, + }, metrics=["correctness"], project=project, ) diff --git a/examples/dataset-experiments/experiment_compare_two_models.py b/examples/dataset-experiments/experiment_compare_two_models.py index 46ec4de3..319c429e 100644 --- a/examples/dataset-experiments/experiment_compare_two_models.py +++ b/examples/dataset-experiments/experiment_compare_two_models.py @@ -2,21 +2,19 @@ This module provides functionality to compare the performance of two LLMs using Splunk AO experiments. """ -import argparse -import json import os +import json import sys +import argparse +from typing import List, Dict, Any from datetime import datetime -from typing import Any - -import anthropic from dotenv import load_dotenv +import anthropic from openai import OpenAI - from splunk_ao import Message, MessageRole -from splunk_ao.datasets import create_dataset -from splunk_ao.experiments import run_experiment from splunk_ao.prompts import create_prompt, get_prompt_template +from splunk_ao.experiments import run_experiment +from splunk_ao.datasets import create_dataset load_dotenv() @@ -77,7 +75,7 @@ def __init__(self): Focus on accuracy and financial impact assessment. """ - def read_jsonl_file(self, file_path: str) -> list[dict[str, Any]]: + def read_jsonl_file(self, file_path: str) -> List[Dict[str, Any]]: """ Reads a JSONL file and returns a list of transactions. @@ -86,7 +84,7 @@ def read_jsonl_file(self, file_path: str) -> list[dict[str, Any]]: """ transactions = [] try: - with open(file_path, encoding="utf-8") as file: + with open(file_path, "r", encoding="utf-8") as file: for line_num, line in enumerate(file, 1): line = line.strip() if line: @@ -105,7 +103,7 @@ def read_jsonl_file(self, file_path: str) -> list[dict[str, Any]]: except Exception as e: raise Exception(f"Error reading JSONL file: {e}") from e - def create_galileo_dataset(self, transactions: list[dict[str, Any]], dataset_name: str) -> Any: + def create_galileo_dataset(self, transactions: List[Dict[str, Any]], dataset_name: str) -> Any: """ Creates a Splunk AO dataset from a list of transactions. @@ -118,7 +116,10 @@ def create_galileo_dataset(self, transactions: list[dict[str, Any]], dataset_nam for transaction in transactions: dataset_content.append({"transaction_data": json.dumps(transaction, ensure_ascii=False)}) - dataset = create_dataset(name=dataset_name, content=dataset_content) + dataset = create_dataset( + name=dataset_name, + content=dataset_content, + ) print(f"Created Splunk AO dataset '{dataset_name}' with {len(transactions)} records") return dataset @@ -176,7 +177,7 @@ def anthropic_llm_call(self, input_data: str) -> str: except Exception as e: return f"Error calling Anthropic API: {e}" - def run_model_experiment(self, experiment_name: str, params: dict[str, Any]) -> None: + def run_model_experiment(self, experiment_name: str, params: Dict[str, Any]) -> None: """ Runs a model experiment using the specified parameters. @@ -188,6 +189,7 @@ def run_model_experiment(self, experiment_name: str, params: dict[str, Any]) -> experiment_name = f"{experiment_name}_{timestamp}" print(f"Running experiment: {experiment_name}") try: + # 2 types of experiments # Runner function experiment: entry point to the app or a function in the codebase. # Prompt experiment @@ -195,11 +197,7 @@ def run_model_experiment(self, experiment_name: str, params: dict[str, Any]) -> experiment_name=experiment_name, dataset=params["dataset"], prompt_template=params["prompt_template"], - prompt_settings={ - "max_tokens": 1000, - "model_alias": params["model_config"]["name"], - "temperature": 0.8, - }, + prompt_settings={"max_tokens": 1000, "model_alias": params["model_config"]["name"], "temperature": 0.8}, metrics=["correctness", "structural_correctness_fin_tx"], project=self.galileo_project, ) @@ -214,11 +212,11 @@ def run_model_experiment(self, experiment_name: str, params: dict[str, Any]) -> except Exception as e: print(f"Error running experiment '{experiment_name}': {e}") - print(f"Full error details: {type(e).__name__}: {e!s}") + print(f"Full error details: {type(e).__name__}: {str(e)}") print(f"Model config: {params['model_config']}") print(f"LLM function: {params['llm_function'].__name__}") - def run_comparison_experiments(self, jsonl_file_path: str, dataset_name: str | None = None) -> None: + def run_comparison_experiments(self, jsonl_file_path: str, dataset_name: str = None) -> None: """ Runs comparison experiments using the provided JSONL file and optional dataset name. @@ -287,17 +285,16 @@ def get_or_create_prompt_template(self) -> Any: return prompt_template except Exception: print(f"Creating new prompt template: {prompt_name}") - return create_prompt( + prompt_template = create_prompt( name=prompt_name, template=[ Message(role=MessageRole.system, content=self.system_prompt), - Message( - role=MessageRole.user, content="Process this financial transaction data: {{transaction_data}}" - ), + Message(role=MessageRole.user, content="Process this financial transaction data: {{transaction_data}}"), ], ) + return prompt_template - def get_optimal_model(self, context: dict[str, Any]) -> str: + def get_optimal_model(self, context: Dict[str, Any]) -> str: """ Determines the optimal model based on the provided context. @@ -311,7 +308,7 @@ def get_optimal_model(self, context: dict[str, Any]) -> str: return max(self.model_configs.keys(), key=lambda x: self.model_configs[x]["performance_score"]) -def main() -> None: +def main(): """ Main function to execute the intelligent broker system for financial data quality. """ @@ -319,11 +316,7 @@ def main() -> None: parser.add_argument("jsonl_file", help="Path to the JSONL file containing financial transactions") parser.add_argument("--dataset-name", help="Optional name for the Splunk AO dataset") - parser.add_argument( - "--project", - default=os.getenv("SPLUNK_AO_PROJECT"), - help="Splunk AO project name (defaults to SPLUNK_AO_PROJECT env var)", - ) + parser.add_argument("--project", default=os.getenv("SPLUNK_AO_PROJECT"), help="Splunk AO project name (defaults to SPLUNK_AO_PROJECT env var)") args = parser.parse_args() diff --git a/examples/experiments/rag-and-tools/app.py b/examples/experiments/rag-and-tools/app.py index d837935d..18674365 100644 --- a/examples/experiments/rag-and-tools/app.py +++ b/examples/experiments/rag-and-tools/app.py @@ -1,10 +1,11 @@ import json -from dotenv import load_dotenv - from splunk_ao import log, splunk_ao_context + from splunk_ao.openai import OpenAI +from dotenv import load_dotenv + # Load environment variables from .env file load_dotenv(override=True) @@ -25,7 +26,10 @@ def retrieve_horoscope_data(sign): "Next Tuesday you will find a four-leaf clover.", "Next Tuesday you will have a great conversation with a stranger.", ], - "Gemini": ["Next Tuesday you will learn to juggle.", "Next Tuesday you will discover a new favorite book."], + "Gemini": [ + "Next Tuesday you will learn to juggle.", + "Next Tuesday you will discover a new favorite book.", + ], } return horoscopes.get(sign, ["No horoscope available."]) @@ -48,12 +52,15 @@ def get_horoscope(sign): "parameters": { "type": "object", "properties": { - "sign": {"type": "string", "description": "An astrological sign like Taurus or Aquarius"} + "sign": { + "type": "string", + "description": "An astrological sign like Taurus or Aquarius", + }, }, "required": ["sign"], }, }, - } + }, ] # Map tool names to their implementations @@ -68,7 +75,11 @@ def call_llm(messages): """ Call the LLM with the provided messages and tools. """ - return client.chat.completions.create(model="gpt-5.1", tools=tools, messages=messages) + return client.chat.completions.create( + model="gpt-5.1", + tools=tools, + messages=messages, + ) def get_users_horoscope(sign: str) -> str: @@ -104,7 +115,10 @@ def get_users_horoscope(sign: str) -> str: { "id": call.id, "type": "function", - "function": {"name": call.function.name, "arguments": call.function.arguments}, + "function": { + "name": call.function.name, + "arguments": call.function.arguments, + }, } for call in completion_tool_calls ], @@ -121,7 +135,12 @@ def get_users_horoscope(sign: str) -> str: # Add the tool result to the message history message_history.append( - {"role": "tool", "content": tool_result, "tool_call_id": call.id, "name": call.function.name} + { + "role": "tool", + "content": tool_result, + "tool_call_id": call.id, + "name": call.function.name, + } ) # Now we call the model again, with the tool results included @@ -131,7 +150,7 @@ def get_users_horoscope(sign: str) -> str: return response.choices[0].message.content -def main() -> None: +def main(): """ Get the user's horoscope """ diff --git a/examples/experiments/rag-and-tools/experiment.py b/examples/experiments/rag-and-tools/experiment.py index 03f2daf0..aa4c19a8 100644 --- a/examples/experiments/rag-and-tools/experiment.py +++ b/examples/experiments/rag-and-tools/experiment.py @@ -1,18 +1,23 @@ import os -from app import get_users_horoscope - from splunk_ao import SplunkAOMetrics from splunk_ao.experiments import run_experiment +from app import get_users_horoscope + -def main() -> None: +def main(): """ Run the horoscope experiment """ # Define a dataset of astrological signs to use # in the experiment - dataset = [{"input": "Aquarius"}, {"input": "Taurus"}, {"input": "Gemini"}, {"input": "Leo"}] + dataset = [ + {"input": "Aquarius"}, + {"input": "Taurus"}, + {"input": "Gemini"}, + {"input": "Leo"}, + ] # Run the experiment results = run_experiment( diff --git a/examples/experiments/upload_experiment/upload_existing_results.py b/examples/experiments/upload_experiment/upload_existing_results.py index f6776d1c..b51f8175 100644 --- a/examples/experiments/upload_experiment/upload_existing_results.py +++ b/examples/experiments/upload_experiment/upload_existing_results.py @@ -21,23 +21,22 @@ - Result: Complete evaluation with metrics and detailed traces """ -import json import os -from typing import Any - +import json +from typing import Dict, Any, Optional from dotenv import load_dotenv # Load environment variables from .env file load_dotenv() # Splunk AO imports -from splunk_ao import splunk_ao_context from splunk_ao.datasets import create_dataset, get_dataset from splunk_ao.experiments import run_experiment from splunk_ao.schema.metrics import SplunkAOMetrics +from splunk_ao import splunk_ao_context -def load_evaluation_data(json_path: str) -> dict[str, dict[str, Any]]: +def load_evaluation_data(json_path: str) -> Dict[str, Dict[str, Any]]: """ Load your existing evaluation results from JSON. @@ -61,7 +60,7 @@ def load_evaluation_data(json_path: str) -> dict[str, dict[str, Any]]: Returns: Dict mapping questions to their full evaluation records """ - with open(json_path) as f: + with open(json_path, "r") as f: data = json.load(f) # Create lookup dict keyed by question @@ -76,11 +75,7 @@ def load_evaluation_data(json_path: str) -> dict[str, dict[str, Any]]: elif not isinstance(context, list): context = [str(context)] if context else [] - lookup[question] = { - "context": context, - "llm_answer": record.get("llm_answer", ""), - "model": record.get("model", "gpt-4o"), - } # Allow model override + lookup[question] = {"context": context, "llm_answer": record.get("llm_answer", ""), "model": record.get("model", "gpt-4o")} # Allow model override return lookup @@ -131,7 +126,7 @@ def prepare_dataset_for_galileo(json_path: str, dataset_name: str) -> Any: Dataset object from Splunk AO """ # Load your evaluation results - with open(json_path) as f: + with open(json_path, "r") as f: raw_data = json.load(f) # Transform to Splunk AO dataset format @@ -144,7 +139,7 @@ def prepare_dataset_for_galileo(json_path: str, dataset_name: str) -> Any: return create_or_get_dataset(dataset_name, galileo_dataset) -def create_replay_function(evaluation_lookup: dict[str, dict[str, Any]], system_prompt: str | None = None): +def create_replay_function(evaluation_lookup: Dict[str, Dict[str, Any]], system_prompt: Optional[str] = None): """ Create a function that replays your evaluation with full tracing. @@ -161,9 +156,7 @@ def create_replay_function(evaluation_lookup: dict[str, dict[str, Any]], system_ # Default system prompt if none provided if system_prompt is None: - system_prompt = ( - "You are a helpful AI assistant. Use the provided context to answer the question accurately and concisely." - ) + system_prompt = "You are a helpful AI assistant. Use the provided context " "to answer the question accurately and concisely." def replay_evaluation(input: str, **kwargs) -> str: """ @@ -196,7 +189,7 @@ def replay_evaluation(input: str, **kwargs) -> str: # Format context chunks for the prompt # Join multiple chunks with clear separators if context_chunks: - context_text = "\n\n---\n\n".join([f"Chunk {i + 1}:\n{chunk}" for i, chunk in enumerate(context_chunks)]) + context_text = "\n\n---\n\n".join([f"Chunk {i+1}:\n{chunk}" for i, chunk in enumerate(context_chunks)]) else: context_text = "" @@ -215,12 +208,7 @@ def replay_evaluation(input: str, **kwargs) -> str: def upload_experiment( - dataset: Any, - evaluation_data_path: str, - project_name: str, - run_name: str, - system_prompt: str | None = None, - metrics: list | None = None, + dataset: Any, evaluation_data_path: str, project_name: str, run_name: str, system_prompt: Optional[str] = None, metrics: Optional[list] = None ) -> Any: """ Upload your evaluation results as a Splunk AO experiment. @@ -259,14 +247,20 @@ def upload_experiment( ] # Run experiment with your data - results = run_experiment(run_name, project=project_name, dataset=dataset, function=replay_fn, metrics=metrics) + results = run_experiment( + run_name, + project=project_name, + dataset=dataset, + function=replay_fn, + metrics=metrics, + ) print("✓ Experiment complete!") return results -def main() -> None: +def main(): """ Example: Upload existing evaluation results to Splunk AO diff --git a/examples/logging-samples/distributed-tracing-otel-python-java/python-service/app.py b/examples/logging-samples/distributed-tracing-otel-python-java/python-service/app.py index 37ae04e9..bad8a4ba 100644 --- a/examples/logging-samples/distributed-tracing-otel-python-java/python-service/app.py +++ b/examples/logging-samples/distributed-tracing-otel-python-java/python-service/app.py @@ -6,6 +6,9 @@ import chromadb.utils.embedding_functions as ef import openai from fastapi import FastAPI +from splunk_ao.otel import start_galileo_span +from splunk_ao_core.schemas.logging.span import RetrieverSpan +from splunk_ao_core.schemas.shared.document import Document from langgraph.graph import END, START, StateGraph from openinference.instrumentation.langchain import LangChainInstrumentor from openinference.instrumentation.openai import OpenAIInstrumentor @@ -16,10 +19,6 @@ from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor from pydantic import BaseModel -from splunk_ao_core.schemas.logging.span import RetrieverSpan -from splunk_ao_core.schemas.shared.document import Document - -from splunk_ao.otel import start_galileo_span logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -71,8 +70,14 @@ # Clients # --------------------------------------------------------------------------- oai = openai.OpenAI() -chroma = chromadb.HttpClient(host=os.getenv("CHROMADB_HOST", "localhost"), port=int(os.getenv("CHROMADB_PORT", "8000"))) -embedding_fn = ef.OpenAIEmbeddingFunction(api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small") +chroma = chromadb.HttpClient( + host=os.getenv("CHROMADB_HOST", "localhost"), + port=int(os.getenv("CHROMADB_PORT", "8000")), +) +embedding_fn = ef.OpenAIEmbeddingFunction( + api_key=os.getenv("OPENAI_API_KEY"), + model_name="text-embedding-3-small", +) CHROMA_COLLECTION = "financial_docs" @@ -115,7 +120,7 @@ def retrieve_documents(state: GraphState) -> GraphState: logger.exception("ChromaDB query failed, continuing with empty context") docs, ids = [], [] - retriever_span.output = [Document(content=d, metadata={"id": i}) for d, i in zip(docs, ids, strict=False)] + retriever_span.output = [Document(content=d, metadata={"id": i}) for d, i in zip(docs, ids)] span.set_attribute("retriever.result_count", len(docs)) logger.info("Retrieved %d documents", len(docs)) @@ -141,10 +146,17 @@ def generate_answer(state: GraphState) -> GraphState: "If the context doesn't contain enough information, say so." ), }, - {"role": "user", "content": f"Context:\n{context_block}\n\nQuestion: {question}"}, + { + "role": "user", + "content": f"Context:\n{context_block}\n\nQuestion: {question}", + }, ] - response = oai.chat.completions.create(model="gpt-4o-mini", messages=messages, temperature=0.2) + response = oai.chat.completions.create( + model="gpt-4o-mini", + messages=messages, + temperature=0.2, + ) # message.content is Optional[str] — guard against content filters or # tool-call responses where the model returns no text. answer = response.choices[0].message.content or "" @@ -195,7 +207,13 @@ def process(req: ProcessRequest): serves as the workflow container in the Splunk AO trace UI. """ result = workflow.invoke( - {"question": req.question, "category": req.category, "documents": [], "sources": [], "answer": ""} + { + "question": req.question, + "category": req.category, + "documents": [], + "sources": [], + "answer": "", + } ) return {"answer": result["answer"], "sources": result.get("sources", [])} diff --git a/examples/logging-samples/distributed-tracing-otel-python-java/python-service/seed_chroma.py b/examples/logging-samples/distributed-tracing-otel-python-java/python-service/seed_chroma.py index fa9b421e..c5813290 100644 --- a/examples/logging-samples/distributed-tracing-otel-python-java/python-service/seed_chroma.py +++ b/examples/logging-samples/distributed-tracing-otel-python-java/python-service/seed_chroma.py @@ -1,6 +1,5 @@ """Seed ChromaDB with sample financial services documents.""" -import contextlib import os import sys import time @@ -76,7 +75,7 @@ ] -def seed() -> None: +def seed(): host = os.getenv("CHROMADB_HOST", "localhost") port = int(os.getenv("CHROMADB_PORT", "8000")) @@ -96,12 +95,20 @@ def seed() -> None: sys.exit(1) # Delete existing collection if present, then create fresh - with contextlib.suppress(Exception): + try: client.delete_collection(COLLECTION) + except Exception: + pass - embedding_fn = ef.OpenAIEmbeddingFunction(api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small") + embedding_fn = ef.OpenAIEmbeddingFunction( + api_key=os.getenv("OPENAI_API_KEY"), + model_name="text-embedding-3-small", + ) collection = client.create_collection(name=COLLECTION, embedding_function=embedding_fn) - collection.add(ids=[d["id"] for d in DOCUMENTS], documents=[d["text"] for d in DOCUMENTS]) + collection.add( + ids=[d["id"] for d in DOCUMENTS], + documents=[d["text"] for d in DOCUMENTS], + ) print(f"Seeded {len(DOCUMENTS)} documents into '{COLLECTION}' collection.") diff --git a/examples/logging-samples/distributed-tracing/main_run.py b/examples/logging-samples/distributed-tracing/main_run.py index e35c53d5..830e781a 100644 --- a/examples/logging-samples/distributed-tracing/main_run.py +++ b/examples/logging-samples/distributed-tracing/main_run.py @@ -11,12 +11,11 @@ Then run this: python main_distributed_tracing.py """ -import asyncio - -import httpx from dotenv import load_dotenv - -from splunk_ao import get_tracing_headers, log, openai +import httpx +import asyncio +from splunk_ao import log, get_tracing_headers +from splunk_ao import openai load_dotenv() @@ -95,13 +94,15 @@ def analyze_question(question: str) -> dict: """ question_lower = question.lower() - return { + analysis = { "needs_company_info": any(word in question_lower for word in ["company", "work", "employer"]), "needs_location_info": any(word in question_lower for word in ["location", "where", "city", "live"]), "needs_education_info": any(word in question_lower for word in ["education", "school", "degree", "study"]), "question_type": "factual", } + return analysis + @log def format_context(analysis: dict, docs: list[str]) -> str: @@ -117,15 +118,18 @@ def format_context(analysis: dict, docs: list[str]) -> str: # ============================================================================ -async def main() -> None: +async def main(): """Run the distributed tracing example""" - questions = ["What did Galileo Galilei research?", "Where did Galileo Galilei work?"] + questions = [ + "What did Galileo Galilei research?", + "Where did Galileo Galilei work?", + ] for question in questions: - print(f"\n{'=' * 60}") + print(f"\n{'='*60}") print(f"Question: {question}") - print(f"{'=' * 60}") + print(f"{'='*60}") answer = await orchestrator_agent(question) print(f"Answer: {answer}\n") diff --git a/examples/logging-samples/distributed-tracing/retrieval_service.py b/examples/logging-samples/distributed-tracing/retrieval_service.py index e5fe3389..7d5a0601 100644 --- a/examples/logging-samples/distributed-tracing/retrieval_service.py +++ b/examples/logging-samples/distributed-tracing/retrieval_service.py @@ -4,10 +4,9 @@ Run this service with: uvicorn retrieval_service:app --reload --port 8000 """ -from dotenv import load_dotenv from fastapi import FastAPI from pydantic import BaseModel - +from dotenv import load_dotenv from splunk_ao import log from splunk_ao.middleware.tracing import TracingMiddleware @@ -37,9 +36,7 @@ def retrieval_service(query: str) -> list[str]: # Mock knowledge base knowledge_base = { "birthplace": ["Galileo Galilei was born in Pisa, Italy in 1564"], - "profession": [ - "Galileo Galilei taught geometry, mechanics, and astronomy at the University of Padua for many years" - ], + "profession": ["Galileo Galilei taught geometry, mechanics, and astronomy at the University of Padua for many years"], "research": [ "Using improved telescopes that he built, Galileo Galilei made scientific observations that transformed our understanding of the universe." ], diff --git a/examples/logging-samples/galileologger/basic-example.py b/examples/logging-samples/galileologger/basic-example.py index 894bbf47..e2a7ed82 100644 --- a/examples/logging-samples/galileologger/basic-example.py +++ b/examples/logging-samples/galileologger/basic-example.py @@ -1,9 +1,9 @@ -# Load environment variables from .env file -from dotenv import load_dotenv - from splunk_ao import SplunkAOLogger from splunk_ao.config import SplunkAOConfig # For displaying the log stream URL +# Load environment variables from .env file +from dotenv import load_dotenv + load_dotenv() # Create a logger instance diff --git a/examples/logging-samples/galileologger/metadata-example.py b/examples/logging-samples/galileologger/metadata-example.py index c477cfd8..f382799c 100644 --- a/examples/logging-samples/galileologger/metadata-example.py +++ b/examples/logging-samples/galileologger/metadata-example.py @@ -1,7 +1,6 @@ -from dotenv import load_dotenv - from splunk_ao import SplunkAOLogger from splunk_ao.config import SplunkAOConfig # For displaying the log stream URL +from dotenv import load_dotenv # Load environment variables from the .env file load_dotenv() @@ -23,9 +22,7 @@ trace = logger.start_trace(input=input_prompt, metadata=metadata) # Add an LLM span to the trace with optional metadata -logger.add_llm_span( - input=[{"role": "system", "content": input_prompt}], output=output_answer, model="gpt-5-mini", metadata=metadata -) +logger.add_llm_span(input=[{"role": "system", "content": input_prompt}], output=output_answer, model="gpt-5-mini", metadata=metadata) # Conclude the trace with the final output logger.conclude(output_answer) diff --git a/examples/logging-samples/galileologger/retriever-example.py b/examples/logging-samples/galileologger/retriever-example.py index c215e5a6..53830a74 100644 --- a/examples/logging-samples/galileologger/retriever-example.py +++ b/examples/logging-samples/galileologger/retriever-example.py @@ -1,9 +1,9 @@ -# Load environment variables from .env file -from dotenv import load_dotenv - from splunk_ao import splunk_ao_context # The Splunk AO context manager from splunk_ao.config import SplunkAOConfig # For displaying the log stream URL +# Load environment variables from .env file +from dotenv import load_dotenv + load_dotenv() prompt_input_data = [ @@ -21,6 +21,7 @@ logger = splunk_ao_context.get_logger_instance() for i in range(len(prompt_input_data)): + prompt_output = prompt_output_data[0] # Start a session @@ -30,12 +31,7 @@ trace = logger.start_trace("This is a trace start") # Add a retriever span for document retrieval - logger.add_retriever_span( - input="Who's a good bot?", - output=["Research shows that I am a good bot."], - name="Document Retrieval", - duration_ns=1000000, - ) + logger.add_retriever_span(input="Who's a good bot?", output=["Research shows that I am a good bot."], name="Document Retrieval", duration_ns=1000000) # Add an LLM span for generating a response logger.add_llm_span( diff --git a/examples/logging-samples/log-mcp-calls/app.py b/examples/logging-samples/log-mcp-calls/app.py index aa79824f..40ddefd2 100644 --- a/examples/logging-samples/log-mcp-calls/app.py +++ b/examples/logging-samples/log-mcp-calls/app.py @@ -7,15 +7,18 @@ import asyncio import os + from datetime import datetime from anthropic import Anthropic, omit from anthropic.types import Message + from dotenv import load_dotenv -from mcp_client import MCPClient from splunk_ao import splunk_ao_context +from mcp_client import MCPClient + load_dotenv() # load environment variables from .env anthropic = Anthropic() @@ -63,7 +66,10 @@ async def process_query(query: str) -> str: # Start a Splunk AO Logger trace galileo_logger = splunk_ao_context.get_logger_instance() - galileo_logger.start_trace(input=query, name="MCP Chatbot Query") + galileo_logger.start_trace( + input=query, + name="MCP Chatbot Query", + ) message_history.append({"role": "user", "content": query}) @@ -113,7 +119,8 @@ async def process_query(query: str) -> str: # Conclude and flush the trace galileo_logger.conclude( - output="\n".join(final_text), duration_ns=int((datetime.now().timestamp() * 1_000_000_000) - start_time_ns) + output="\n".join(final_text), + duration_ns=int((datetime.now().timestamp() * 1_000_000_000) - start_time_ns), ) galileo_logger.flush() @@ -121,7 +128,7 @@ async def process_query(query: str) -> str: return "\n".join(final_text) -async def main() -> None: +async def main(): """Main function to run the chat loop""" # Connect to the MCP server await mcp_client.connect_to_server() diff --git a/examples/logging-samples/log-mcp-calls/mcp_client.py b/examples/logging-samples/log-mcp-calls/mcp_client.py index 8f2f97a9..4236ac6a 100644 --- a/examples/logging-samples/log-mcp-calls/mcp_client.py +++ b/examples/logging-samples/log-mcp-calls/mcp_client.py @@ -3,7 +3,9 @@ """ import os + from contextlib import AsyncExitStack +from typing import Optional from mcp import ClientSession from mcp.client.streamable_http import streamablehttp_client @@ -14,17 +16,20 @@ class MCPClient: def __init__(self): # Initialize session and client objects - self._session: ClientSession | None = None + self._session: Optional[ClientSession] = None self._exit_stack = AsyncExitStack() self.tools = [] - async def connect_to_server(self) -> None: + async def connect_to_server(self): """Connect to an MCP server""" # Establish streamable HTTP connection read, write, _ = await self._exit_stack.enter_async_context( streamablehttp_client( url=os.environ.get("MCP_SERVER_URL", "https://api.galileo.ai/mcp/http/mcp"), - headers={"Splunk-AO-API-Key": os.environ["SPLUNK_AO_API_KEY"], "Accept": "text/event-stream"}, + headers={ + "Splunk-AO-API-Key": os.environ["SPLUNK_AO_API_KEY"], + "Accept": "text/event-stream", + }, ) ) @@ -37,10 +42,17 @@ async def connect_to_server(self) -> None: # List the available tools response = await self._session.list_tools() self.tools = [ - {"name": tool.name, "description": tool.description, "input_schema": tool.inputSchema} + { + "name": tool.name, + "description": tool.description, + "input_schema": tool.inputSchema, + } for tool in response.tools ] - print("\nConnected to server with tools:", [tool["name"] for tool in self.tools]) + print( + "\nConnected to server with tools:", + [tool["name"] for tool in self.tools], + ) async def call_tool(self, tool_name: str, input_data: dict): """Call a tool by name with the provided input data""" @@ -51,6 +63,6 @@ async def call_tool(self, tool_name: str, input_data: dict): # Call the tool and return the result return await self._session.call_tool(tool_name, input_data) - async def cleanup(self) -> None: + async def cleanup(self): """Clean up resources""" await self._exit_stack.aclose() diff --git a/examples/logging-samples/openai-responses/main.py b/examples/logging-samples/openai-responses/main.py index e7d6b011..e523efbf 100644 --- a/examples/logging-samples/openai-responses/main.py +++ b/examples/logging-samples/openai-responses/main.py @@ -26,8 +26,15 @@ "parameters": { "type": "object", "properties": { - "location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"}, - "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "The temperature unit"}, + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + "description": "The temperature unit", + }, }, "required": ["location"], }, @@ -39,7 +46,10 @@ "parameters": { "type": "object", "properties": { - "symbol": {"type": "string", "description": "The stock ticker symbol, e.g. AAPL, GOOGL, MSFT"} + "symbol": { + "type": "string", + "description": "The stock ticker symbol, e.g. AAPL, GOOGL, MSFT", + }, }, "required": ["symbol"], }, @@ -63,7 +73,7 @@ def get_stock_price(symbol: str) -> str: return json.dumps({"symbol": symbol.upper(), "price": price, "currency": "USD"}) -def main() -> None: +def main(): user_message = "What's the weather like in San Francisco and what's the current stock price of Apple?" input_list = [] diff --git a/examples/rag/cli-rag-demo/app.py b/examples/rag/cli-rag-demo/app.py index a95fedc0..8217279d 100644 --- a/examples/rag/cli-rag-demo/app.py +++ b/examples/rag/cli-rag-demo/app.py @@ -1,13 +1,11 @@ import os -import sys - -import questionary from dotenv import load_dotenv +from splunk_ao import openai, log, SplunkAOLogger from rich.console import Console -from rich.markdown import Markdown from rich.panel import Panel - -from splunk_ao import SplunkAOLogger, log, openai +from rich.markdown import Markdown +import questionary +import sys load_dotenv() @@ -17,7 +15,10 @@ # Check if Splunk AO logging is enabled logging_enabled = os.environ.get("SPLUNK_AO_API_KEY") is not None -logger = SplunkAOLogger(project="rag-test", log_stream="dev") +logger = SplunkAOLogger( + project="rag-test", + log_stream="dev", +) # Initialize OpenAI client directly client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY")) @@ -26,7 +27,7 @@ @log(span_type="retriever") def retrieve_documents(query: str): # TODO: Replace with actual RAG retrieval - return [ + documents = [ { "id": "doc1", "text": ( @@ -63,6 +64,7 @@ def retrieve_documents(query: str): "metadata": {"source": "best_practices", "category": "prompting"}, }, ] + return documents def rag(query: str): @@ -71,7 +73,7 @@ def rag(query: str): # Format documents for better readability in the prompt formatted_docs = "" for i, doc in enumerate(documents): - formatted_docs += f"Document {i + 1} (Source: {doc['metadata']['source']}):\n{doc['text']}\n\n" + formatted_docs += f"Document {i+1} (Source: {doc['metadata']['source']}):\n{doc['text']}\n\n" prompt = f""" Answer the following question based on the context provided. If the answer is not in the context, say you don't know. @@ -96,10 +98,10 @@ def rag(query: str): ) return response.choices[0].message.content.strip() except Exception as e: - return f"Error generating response: {e!s}" + return f"Error generating response: {str(e)}" -def main() -> None: +def main(): console.print( Panel.fit( "[bold]RAG Demo[/bold]\nThis demo uses a simulated RAG system to answer your questions.", @@ -125,7 +127,8 @@ def main() -> None: while True: # Get user query query = questionary.text( - "Enter your question about Splunk AO, RAG, or AI techniques:", validate=lambda text: len(text) > 0 + "Enter your question about Splunk AO, RAG, or AI techniques:", + validate=lambda text: len(text) > 0, ).ask() if query.lower() in ["exit", "quit", "q"]: @@ -144,7 +147,7 @@ def main() -> None: break except Exception as e: - console.print(f"[bold red]Error:[/bold red] {e!s}") + console.print(f"[bold red]Error:[/bold red] {str(e)}") if __name__ == "__main__": diff --git a/examples/rag/cli-rag-demo/challenges/chunk-utilization.py b/examples/rag/cli-rag-demo/challenges/chunk-utilization.py index 52763d5b..c915d092 100644 --- a/examples/rag/cli-rag-demo/challenges/chunk-utilization.py +++ b/examples/rag/cli-rag-demo/challenges/chunk-utilization.py @@ -1,17 +1,15 @@ import os -import random -import sys -from time import perf_counter_ns - -import faiss -import questionary from dotenv import load_dotenv +from splunk_ao import openai, SplunkAOLogger from rich.console import Console -from rich.markdown import Markdown from rich.panel import Panel +from rich.markdown import Markdown +import questionary +import sys +import faiss from sentence_transformers import SentenceTransformer - -from splunk_ao import SplunkAOLogger, openai +import random +from time import perf_counter_ns load_dotenv() @@ -23,7 +21,10 @@ # Initialize Splunk AO logger -logger = SplunkAOLogger(project="chunk-utilization", log_stream="dev") +logger = SplunkAOLogger( + project="chunk-utilization", + log_stream="dev", +) # Initialize OpenAI client client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY")) @@ -40,7 +41,7 @@ def __init__(self): self._initialize_documents() self._build_index() - def _initialize_documents(self) -> None: + def _initialize_documents(self): # Solar System documents self.documents.extend( [ @@ -60,14 +61,22 @@ def _initialize_documents(self) -> None: "text": ( "The four inner system planets—Mercury, Venus, Earth and Mars—are terrestrial planets, being composed primarily of rock and metal. The four giant planets of the outer system are substantially larger and more massive than the terrestrials. The two largest, Jupiter and Saturn, are gas giants, being composed mainly of hydrogen and helium; the next two, Uranus and Neptune, are ice giants." ), - "metadata": {"source": "astronomy_textbook", "category": "planetary systems", "relevance": "high"}, + "metadata": { + "source": "astronomy_textbook", + "category": "planetary systems", + "relevance": "high", + }, }, { "id": "solar_system_3", "text": ( "The history of Solar System observation dates back to ancient times when astronomers first noticed that certain lights moved across the sky in a different way than the fixed stars. Ancient Greeks called these lights 'planetai' or wanderers, giving rise to our modern term 'planet'." ), - "metadata": {"source": "astronomy_history", "category": "astronomy history", "relevance": "low"}, + "metadata": { + "source": "astronomy_history", + "category": "astronomy history", + "relevance": "low", + }, }, ] ) @@ -80,7 +89,11 @@ def _initialize_documents(self) -> None: "text": ( "Photosynthesis is the process by which plants, algae, and certain bacteria convert light energy, typically from the Sun, into chemical energy in the form of glucose or other sugars. These organisms are called photoautotrophs since they can create their own food." ), - "metadata": {"source": "biology_textbook", "category": "cellular processes", "relevance": "high"}, + "metadata": { + "source": "biology_textbook", + "category": "cellular processes", + "relevance": "high", + }, }, { "id": "photosynthesis_2", @@ -98,7 +111,11 @@ def _initialize_documents(self) -> None: "text": ( "The evolution of photosynthesis occurred early in Earth's history, with the earliest photosynthetic organisms appearing between 3.4 and 2.9 billion years ago. This development dramatically changed Earth's atmosphere by introducing oxygen." ), - "metadata": {"source": "evolutionary_biology", "category": "evolution", "relevance": "low"}, + "metadata": { + "source": "evolutionary_biology", + "category": "evolution", + "relevance": "low", + }, }, ] ) @@ -106,7 +123,7 @@ def _initialize_documents(self) -> None: # Add more topics similarly... # (Blockchain, Renaissance, and Machine Learning documents would be added here) - def _build_index(self) -> None: + def _build_index(self): # Generate embeddings for all documents texts = [doc["text"] for doc in self.documents] self.document_embeddings = encoder.encode(texts) @@ -121,7 +138,7 @@ def search(self, query: str, k: int = 3, mixed_relevance: bool = False) -> list: query_vector = encoder.encode([query])[0].reshape(1, -1) # Search the index - _distances, indices = self.index.search(query_vector.astype("float32"), k) + distances, indices = self.index.search(query_vector.astype("float32"), k) # Get the retrieved documents retrieved_docs = [] @@ -174,7 +191,10 @@ def retrieve_verbose_documents(query: str, mixed_relevance: bool = False): return documents except Exception as e: logger.add_retriever_span( - input=query, output=str(e), duration_ns=perf_counter_ns() - start_time, status_code=500 + input=query, + output=str(e), + duration_ns=perf_counter_ns() - start_time, + status_code=500, ) raise e @@ -190,7 +210,7 @@ def rag_with_poor_utilization(query: str, mixed_relevance: bool = False): # Format documents for the prompt formatted_docs = "" for i, doc in enumerate(documents): - formatted_docs += f"Document {i + 1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}):\n{doc['content']}\n\n" + formatted_docs += f"Document {i+1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}):\n{doc['content']}\n\n" # Basic prompt basic_prompt = f""" @@ -225,9 +245,13 @@ def rag_with_poor_utilization(query: str, mixed_relevance: bool = False): return result except Exception as e: - error_msg = f"Error generating response: {e!s}" + error_msg = f"Error generating response: {str(e)}" logger.add_llm_span( - input=query, output=error_msg, model="gpt-4", duration_ns=perf_counter_ns() - start_time, status_code=500 + input=query, + output=error_msg, + model="gpt-4", + duration_ns=perf_counter_ns() - start_time, + status_code=500, ) return error_msg @@ -243,7 +267,7 @@ def rag_with_better_utilization(query: str, mixed_relevance: bool = False): # Format documents for the prompt formatted_docs = "" for i, doc in enumerate(documents): - formatted_docs += f"Document {i + 1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}):\n{doc['content']}\n\n" + formatted_docs += f"Document {i+1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}):\n{doc['content']}\n\n" # Enhanced prompt enhanced_prompt = f""" @@ -291,14 +315,18 @@ def rag_with_better_utilization(query: str, mixed_relevance: bool = False): return result except Exception as e: - error_msg = f"Error generating response: {e!s}" + error_msg = f"Error generating response: {str(e)}" logger.add_llm_span( - input=query, output=error_msg, model="gpt-4", duration_ns=perf_counter_ns() - start_time, status_code=500 + input=query, + output=error_msg, + model="gpt-4", + duration_ns=perf_counter_ns() - start_time, + status_code=500, ) return error_msg -def main() -> None: +def main(): start_time = perf_counter_ns() try: # Start trace with a meaningful name, actual input will be added when we get the query @@ -337,16 +365,18 @@ def main() -> None: "What are the different types of machine learning?", ] - console.print( - "\n[bold yellow]Suggested queries (these will demonstrate the chunk utilization problem):[/bold yellow]" - ) + console.print("\n[bold yellow]Suggested queries (these will demonstrate the chunk utilization problem):[/bold yellow]") for i, q in enumerate(suggested_queries): - console.print(f"[yellow]{i + 1}. {q}[/yellow]") + console.print(f"[yellow]{i+1}. {q}[/yellow]") # Choose workflow workflow = questionary.select( "Choose which workflow to run:", - choices=["Poor Chunk Utilization (Bad State)", "Improved Chunk Utilization (Good State)", "Exit"], + choices=[ + "Poor Chunk Utilization (Bad State)", + "Improved Chunk Utilization (Good State)", + "Exit", + ], ).ask() if workflow == "Exit": @@ -374,7 +404,8 @@ def main() -> None: # Ask about mixed relevance mixed_relevance = questionary.confirm( - "Would you like to include mixed relevance results (some less relevant documents)?", default=False + "Would you like to include mixed relevance results (some less relevant documents)?", + default=False, ).ask() try: @@ -386,7 +417,7 @@ def main() -> None: console.print( Panel( f"[bold]Source:[/bold] {doc['metadata']['source']}\n[bold]Relevance:[/bold] [{relevance_color}]{doc['metadata']['relevance']}[/{relevance_color}]\n\n[dim]{doc['content']}[/dim]", - title=f"Document {i + 1} ({len(doc['content'])} characters)", + title=f"Document {i+1} ({len(doc['content'])} characters)", border_style="cyan", ) ) @@ -414,12 +445,19 @@ def main() -> None: break except Exception as e: - console.print(f"[bold red]Error:[/bold red] {e!s}") + console.print(f"[bold red]Error:[/bold red] {str(e)}") # Final conclusion of the demo - logger.conclude(output="Demo completed successfully", duration_ns=perf_counter_ns() - start_time) + logger.conclude( + output="Demo completed successfully", + duration_ns=perf_counter_ns() - start_time, + ) except Exception as e: - logger.conclude(output=f"Demo failed: {e!s}", duration_ns=perf_counter_ns() - start_time, status_code=500) + logger.conclude( + output=f"Demo failed: {str(e)}", + duration_ns=perf_counter_ns() - start_time, + status_code=500, + ) raise e diff --git a/examples/rag/cli-rag-demo/challenges/chunk-utilizations-basic.py b/examples/rag/cli-rag-demo/challenges/chunk-utilizations-basic.py index 21367dcb..0230460b 100644 --- a/examples/rag/cli-rag-demo/challenges/chunk-utilizations-basic.py +++ b/examples/rag/cli-rag-demo/challenges/chunk-utilizations-basic.py @@ -1,12 +1,10 @@ import os from pathlib import Path - -from document_store import DocumentStore, format_documents from dotenv import load_dotenv +from splunk_ao.openai import openai from rich.console import Console - +from document_store import DocumentStore, format_documents from splunk_ao import splunk_ao_context -from splunk_ao.openai import openai # Find the .env file in the parent directory current_dir = Path(__file__).resolve().parent @@ -22,9 +20,7 @@ print("SPLUNK_AO_API_KEY:", os.getenv("SPLUNK_AO_API_KEY")) print("SPLUNK_AO_BASE_URL:", os.getenv("SPLUNK_AO_BASE_URL")) -EXAMPLE_QUESTION = ( - "What are the fundamental concepts and operations in arithmetic, and how are they used in mathematics?" -) +EXAMPLE_QUESTION = "What are the fundamental concepts and operations in arithmetic, and how are they used in mathematics?" class Prompts: @@ -64,8 +60,11 @@ def query(question: str): return response.choices[0].message.content.strip() -def main() -> None: - with splunk_ao_context(project=os.getenv("SPLUNK_AO_PROJECT", "chunk-utilization"), log_stream="basic_approach"): +def main(): + with splunk_ao_context( + project=os.getenv("SPLUNK_AO_PROJECT", "chunk-utilization"), + log_stream="basic_approach", + ): console = Console() console.print("\nBasic Chunk Utilization Demo") console.print("\nUsing example question:", EXAMPLE_QUESTION) diff --git a/examples/rag/cli-rag-demo/challenges/chunk-utilizations-enhanced.py b/examples/rag/cli-rag-demo/challenges/chunk-utilizations-enhanced.py index 12b1ac0e..6f0b7587 100644 --- a/examples/rag/cli-rag-demo/challenges/chunk-utilizations-enhanced.py +++ b/examples/rag/cli-rag-demo/challenges/chunk-utilizations-enhanced.py @@ -1,12 +1,10 @@ import os from pathlib import Path - -from document_store import DocumentStore, format_documents from dotenv import load_dotenv +from splunk_ao.openai import openai from rich.console import Console - +from document_store import DocumentStore, format_documents from splunk_ao import splunk_ao_context -from splunk_ao.openai import openai # Find the .env file in the parent directory current_dir = Path(__file__).resolve().parent @@ -16,9 +14,7 @@ # Load the .env file load_dotenv(dotenv_path) -EXAMPLE_QUESTION = ( - "What are the fundamental concepts and operations in arithmetic, and how are they used in mathematics?" -) +EXAMPLE_QUESTION = "What are the fundamental concepts and operations in arithmetic, and how are they used in mathematics?" SYSTEM_PROMPT = """You are a knowledgeable mathematics educator tasked with providing comprehensive answers by analyzing and synthesizing information from multiple provided documents. @@ -43,7 +39,7 @@ class Prompts: def format_documents_enhanced(documents: list) -> str: return "\n\n".join( - f"Document {i + 1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}, " + f"Document {i+1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}, " f"Score: {doc['metadata'].get('combined_score', doc['metadata'].get('score', 'N/A')):.3f}):\n{doc['text']}" for i, doc in enumerate(documents) ) @@ -58,14 +54,21 @@ def query(question: str): prompt = Prompts.ENHANCED.format(query=question, documents=format_documents(docs)) response = client.chat.completions.create( - model="gpt-4", messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}] + model="gpt-4", + messages=[ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], ) return response.choices[0].message.content.strip() -def main() -> None: - with splunk_ao_context(project=os.getenv("SPLUNK_AO_PROJECT", "chunk-utilization"), log_stream="enhanced_approach"): +def main(): + with splunk_ao_context( + project=os.getenv("SPLUNK_AO_PROJECT", "chunk-utilization"), + log_stream="enhanced_approach", + ): console = Console() console.print("\nEnhanced Chunk Utilization Demo") console.print("\nUsing example question:", EXAMPLE_QUESTION) diff --git a/examples/rag/cli-rag-demo/challenges/document_store.py b/examples/rag/cli-rag-demo/challenges/document_store.py index bf85df6f..1229cef8 100644 --- a/examples/rag/cli-rag-demo/challenges/document_store.py +++ b/examples/rag/cli-rag-demo/challenges/document_store.py @@ -1,26 +1,25 @@ -import re -from pathlib import Path - +from sentence_transformers import SentenceTransformer, CrossEncoder import faiss -import numpy as np from datasets import load_dataset -from sentence_transformers import CrossEncoder, SentenceTransformer - +import numpy as np from splunk_ao import log +from typing import List, Dict, Optional +import re +from pathlib import Path class DocumentStore: def __init__( self, source: str = "wikipedia", # "wikipedia" or "custom" - custom_documents_path: str | None = None, + custom_documents_path: Optional[str] = None, num_docs: int = 1000, k: int = 3, chunk_size: int = 512, reranking_threshold: float = 0.6, use_reranking: bool = False, reranking_multiplier: int = 4, - wikipedia_query: str | None = None, + wikipedia_query: Optional[str] = None, ): self.source = source self.documents = [] @@ -41,7 +40,10 @@ def __init__( # Load Wikipedia dataset with configurable parameters print(f"Loading {num_docs} Wikipedia articles...") dataset = load_dataset( - "wikipedia", "20220301.simple", split=f"train[:{num_docs}]", trust_remote_code=True + "wikipedia", + "20220301.simple", + split=f"train[:{num_docs}]", + trust_remote_code=True, ) if wikipedia_query: @@ -88,16 +90,12 @@ def __init__( if not self.documents: raise ValueError("No documents were successfully loaded") - print( - f"Processed {len(self.documents)} chunks from {len({doc['metadata']['source'] for doc in self.documents})} documents" - ) - print( - f"Average chunk length: {sum(doc['metadata']['length'] for doc in self.documents) / len(self.documents):.0f} characters" - ) + print(f"Processed {len(self.documents)} chunks from {len(set(doc['metadata']['source'] for doc in self.documents))} documents") + print(f"Average chunk length: {sum(doc['metadata']['length'] for doc in self.documents) / len(self.documents):.0f} characters") self._build_index() - def _process_wikipedia_documents(self, dataset, chunk_size: int) -> None: + def _process_wikipedia_documents(self, dataset, chunk_size: int): """Process Wikipedia documents into chunks""" for item in dataset: # Skip empty or very short articles @@ -123,13 +121,13 @@ def _process_wikipedia_documents(self, dataset, chunk_size: int) -> None: } ) - def _load_custom_documents(self, documents_path: str, chunk_size: int) -> None: + def _load_custom_documents(self, documents_path: str, chunk_size: int): """Helper method to load custom documents from a file""" documents_path = Path(documents_path) if not documents_path.exists(): raise FileNotFoundError(f"Custom documents file not found: {documents_path}") - with open(documents_path) as f: + with open(documents_path, "r") as f: text = f.read() # Split into paragraphs and process as documents @@ -149,7 +147,7 @@ def _load_custom_documents(self, documents_path: str, chunk_size: int) -> None: { "text": chunk, "metadata": { - "source": f"Document {i + 1}", + "source": f"Document {i+1}", "chunk_id": j, "total_chunks": len(chunks), "relevance": "medium", @@ -160,7 +158,7 @@ def _load_custom_documents(self, documents_path: str, chunk_size: int) -> None: } ) - def _chunk_text(self, text: str, chunk_size: int) -> list[str]: + def _chunk_text(self, text: str, chunk_size: int) -> List[str]: """ Split text into chunks of approximately chunk_size tokens. Uses sentence boundaries where possible to maintain context. @@ -192,7 +190,7 @@ def _chunk_text(self, text: str, chunk_size: int) -> list[str]: return chunks - def _build_index(self) -> None: + def _build_index(self): print("Building FAISS index...") texts = [doc["text"] for doc in self.documents] self.embeddings = self.encoder.encode(texts) @@ -202,13 +200,11 @@ def _build_index(self) -> None: faiss.normalize_L2(self.embeddings) # Create an index that's optimized for cosine similarity - self.index = faiss.IndexFlatIP( - dimension - ) # Inner product is equivalent to cosine similarity for normalized vectors + self.index = faiss.IndexFlatIP(dimension) # Inner product is equivalent to cosine similarity for normalized vectors self.index.add(self.embeddings.astype("float32")) print("Index built successfully") - def _rerank_documents(self, docs: list[dict], query: str) -> list[dict]: + def _rerank_documents(self, docs: List[Dict], query: str) -> List[Dict]: """ Rerank documents using a cross-encoder model, which provides more accurate relevance scoring by considering query and document together. @@ -225,15 +221,13 @@ def _rerank_documents(self, docs: list[dict], query: str) -> list[dict]: score_range = max_score - min_score reranked_docs = [] - for doc, score in zip(docs, cross_scores, strict=False): + for doc, score in zip(docs, cross_scores): # Normalize score to 0-1 range normalized_score = (score - min_score) / score_range if score_range > 0 else 0.5 if normalized_score >= self.reranking_threshold: doc["metadata"]["combined_score"] = normalized_score - doc["metadata"]["relevance"] = ( - "high" if normalized_score > 0.8 else "medium" if normalized_score > 0.6 else "low" - ) + doc["metadata"]["relevance"] = "high" if normalized_score > 0.8 else "medium" if normalized_score > 0.6 else "low" reranked_docs.append(doc) # Sort by combined score @@ -253,7 +247,7 @@ def search(self, query: str) -> list: # Process results results = [] - for score, idx in zip(scores[0], indices[0], strict=False): + for score, idx in zip(scores[0], indices[0]): doc = self.documents[idx].copy() doc["metadata"] = doc["metadata"].copy() doc["metadata"]["score"] = float(score) @@ -263,17 +257,16 @@ def search(self, query: str) -> list: if self.use_reranking: print(f"Reranking top {initial_k} results...") return self._rerank_documents(results, query) - # For basic search, just update relevance based on score - for doc in results: - doc["metadata"]["relevance"] = ( - "high" if doc["metadata"]["score"] > 0.8 else "medium" if doc["metadata"]["score"] > 0.6 else "low" - ) - return results[: self.k] + else: + # For basic search, just update relevance based on score + for doc in results: + doc["metadata"]["relevance"] = "high" if doc["metadata"]["score"] > 0.8 else "medium" if doc["metadata"]["score"] > 0.6 else "low" + return results[: self.k] def format_documents(documents: list) -> str: return "\n\n".join( - f"Document {i + 1} (Source: {doc['metadata']['source']}, Chunk {doc['metadata']['chunk_id'] + 1}/{doc['metadata']['total_chunks']}, " + f"Document {i+1} (Source: {doc['metadata']['source']}, Chunk {doc['metadata']['chunk_id'] + 1}/{doc['metadata']['total_chunks']}, " f"Relevance: {doc['metadata']['relevance']}, " f"Score: {doc['metadata'].get('combined_score', doc['metadata'].get('score', 'N/A')):.3f}):\n{doc['text']}" for i, doc in enumerate(documents) diff --git a/examples/rag/cli-rag-demo/challenges/document_store_basic.py b/examples/rag/cli-rag-demo/challenges/document_store_basic.py index d86a9e88..74cb2dc4 100644 --- a/examples/rag/cli-rag-demo/challenges/document_store_basic.py +++ b/examples/rag/cli-rag-demo/challenges/document_store_basic.py @@ -1,15 +1,15 @@ +from sentence_transformers import SentenceTransformer +import faiss +from typing import Optional import re from pathlib import Path -import faiss -from sentence_transformers import SentenceTransformer - class DocumentStoreBasic: def __init__( self, source: str = "custom", - custom_documents_path: str | None = None, + custom_documents_path: Optional[str] = None, num_docs: int = 1, # Only return 1 document chunk_size: int = 1000, # Larger chunks, less precise ): @@ -35,13 +35,13 @@ def __init__( self._build_index() - def _load_custom_documents(self, documents_path: str, chunk_size: int) -> None: + def _load_custom_documents(self, documents_path: str, chunk_size: int): """Helper method to load custom documents from a file""" documents_path = Path(documents_path) if not documents_path.exists(): raise FileNotFoundError(f"Custom documents file not found: {documents_path}") - with open(documents_path) as f: + with open(documents_path, "r") as f: text = f.read() # Simple chunking by paragraphs @@ -75,7 +75,7 @@ def _load_custom_documents(self, documents_path: str, chunk_size: int) -> None: { "text": chunk, "metadata": { - "source": f"Document {i + 1}", + "source": f"Document {i+1}", "chunk_id": j, "total_chunks": len(chunks), "relevance": "medium", @@ -84,7 +84,7 @@ def _load_custom_documents(self, documents_path: str, chunk_size: int) -> None: } ) - def _build_index(self) -> None: + def _build_index(self): print("Building basic FAISS index...") texts = [doc["text"] for doc in self.documents] self.embeddings = self.encoder.encode(texts) @@ -105,7 +105,7 @@ def search(self, query: str) -> list: # Process results results = [] - for distance, idx in zip(distances[0], indices[0], strict=False): + for distance, idx in zip(distances[0], indices[0]): doc = self.documents[idx].copy() doc["metadata"] = doc["metadata"].copy() doc["metadata"]["score"] = float(1 / (1 + distance)) # Simple distance to score conversion @@ -115,4 +115,4 @@ def search(self, query: str) -> list: def format_documents(documents: list) -> str: - return "\n\n".join(f"Document {i + 1}:\n{doc['text']}" for i, doc in enumerate(documents)) + return "\n\n".join(f"Document {i+1}:\n{doc['text']}" for i, doc in enumerate(documents)) diff --git a/examples/rag/cli-rag-demo/challenges/ensure-completeness-basic.py b/examples/rag/cli-rag-demo/challenges/ensure-completeness-basic.py index 0be10a68..c40a04d8 100644 --- a/examples/rag/cli-rag-demo/challenges/ensure-completeness-basic.py +++ b/examples/rag/cli-rag-demo/challenges/ensure-completeness-basic.py @@ -1,12 +1,10 @@ import os from pathlib import Path - -from document_store_basic import DocumentStoreBasic, format_documents from dotenv import load_dotenv +from splunk_ao.openai import openai from rich.console import Console - +from document_store_basic import DocumentStoreBasic, format_documents from splunk_ao import splunk_ao_context -from splunk_ao.openai import openai # Find the .env file in the parent directory current_dir = Path(__file__).resolve().parent @@ -47,14 +45,21 @@ def query(question: str): prompt = Prompts.BASIC.format(query=question, documents=format_documents(docs)) response = client.chat.completions.create( - model="gpt-4", messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}] + model="gpt-4", + messages=[ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], ) return response.choices[0].message.content.strip() -def main() -> None: - with splunk_ao_context(project=os.getenv("SPLUNK_AO_PROJECT", "ensure-completeness"), log_stream="basic_approach"): +def main(): + with splunk_ao_context( + project=os.getenv("SPLUNK_AO_PROJECT", "ensure-completeness"), + log_stream="basic_approach", + ): console = Console() console.print("\nBasic Completeness Demo") console.print("\nUsing example question:", EXAMPLE_QUESTION) diff --git a/examples/rag/cli-rag-demo/challenges/ensure-completeness-enhanced.py b/examples/rag/cli-rag-demo/challenges/ensure-completeness-enhanced.py index 84e08a63..b960561d 100644 --- a/examples/rag/cli-rag-demo/challenges/ensure-completeness-enhanced.py +++ b/examples/rag/cli-rag-demo/challenges/ensure-completeness-enhanced.py @@ -1,12 +1,10 @@ import os from pathlib import Path - -from document_store import DocumentStore from dotenv import load_dotenv +from splunk_ao.openai import openai from rich.console import Console - +from document_store import DocumentStore from splunk_ao import splunk_ao_context -from splunk_ao.openai import openai # Find the .env file in the parent directory current_dir = Path(__file__).resolve().parent @@ -47,7 +45,7 @@ class Prompts: def format_documents_with_citations(documents: list) -> str: return "\n\n".join( - f"Document {i + 1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}, " + f"Document {i+1} (Source: {doc['metadata']['source']}, Relevance: {doc['metadata']['relevance']}, " f"Score: {doc['metadata'].get('combined_score', doc['metadata'].get('score', 'N/A')):.3f}):\n{doc['text']}" for i, doc in enumerate(documents) ) @@ -72,15 +70,20 @@ def query(question: str): prompt = Prompts.ENHANCED.format(query=question, documents=format_documents_with_citations(docs)) response = client.chat.completions.create( - model="gpt-4", messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}] + model="gpt-4", + messages=[ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], ) return response.choices[0].message.content.strip() -def main() -> None: +def main(): with splunk_ao_context( - project=os.getenv("SPLUNK_AO_PROJECT", "ensure-completeness"), log_stream="enhanced_approach" + project=os.getenv("SPLUNK_AO_PROJECT", "ensure-completeness"), + log_stream="enhanced_approach", ): console = Console() console.print("\nEnhanced Completeness Demo") diff --git a/examples/rag/cli-rag-demo/challenges/out-of-context.py b/examples/rag/cli-rag-demo/challenges/out-of-context.py index 8037e13b..b726bc5d 100644 --- a/examples/rag/cli-rag-demo/challenges/out-of-context.py +++ b/examples/rag/cli-rag-demo/challenges/out-of-context.py @@ -1,9 +1,7 @@ import os - -import questionary from dotenv import load_dotenv - -from splunk_ao import log, openai, splunk_ao_context +from splunk_ao import openai, log, splunk_ao_context +import questionary load_dotenv() @@ -27,13 +25,23 @@ def retrieve_documents(query: str): "eiffel tower": [ { "content": "The Eiffel Tower is an iron lattice tower located in Paris, France. It was designed by Gustave Eiffel.", - "metadata": {"id": "doc1", "source": "travel_guide", "category": "landmarks", "relevance": "high"}, + "metadata": { + "id": "doc1", + "source": "travel_guide", + "category": "landmarks", + "relevance": "high", + }, } ], "python language": [ { "content": "Python is a high-level programming language known for its readability and simple syntax.", - "metadata": {"id": "doc1", "source": "programming_guide", "category": "languages", "relevance": "high"}, + "metadata": { + "id": "doc1", + "source": "programming_guide", + "category": "languages", + "relevance": "high", + }, } ], "climate change": [ @@ -52,7 +60,12 @@ def retrieve_documents(query: str): "artificial intelligence": [ { "content": "Artificial intelligence involves creating systems capable of performing tasks that typically require human intelligence.", - "metadata": {"id": "doc1", "source": "technology_overview", "category": "ai", "relevance": "high"}, + "metadata": { + "id": "doc1", + "source": "technology_overview", + "category": "ai", + "relevance": "high", + }, } ], "quantum computing": [ @@ -100,7 +113,7 @@ def rag_with_hallucination(query: str): # Format documents for better readability in the prompt formatted_docs = "" for i, doc in enumerate(documents): - formatted_docs += f"Document {i + 1} (Source: {doc['metadata']['source']}):\n{doc['content']}\n\n" + formatted_docs += f"Document {i+1} (Source: {doc['metadata']['source']}):\n{doc['content']}\n\n" # This prompt doesn't strongly constrain the model to only use the provided context weak_prompt = f""" @@ -123,7 +136,7 @@ def rag_with_hallucination(query: str): ) return response.choices[0].message.content.strip() except Exception as e: - return f"Error generating response: {e!s}" + return f"Error generating response: {str(e)}" @log(name="rag_with_constraint") @@ -137,7 +150,7 @@ def rag_with_constraint(query: str): # Format documents for better readability in the prompt formatted_docs = "" for i, doc in enumerate(documents): - formatted_docs += f"Document {i + 1} (Source: {doc['metadata']['source']}):\n{doc['content']}\n\n" + formatted_docs += f"Document {i+1} (Source: {doc['metadata']['source']}):\n{doc['content']}\n\n" # This prompt strongly constrains the model to only use the provided context strong_prompt = f""" @@ -167,11 +180,11 @@ def rag_with_constraint(query: str): ) return response.choices[0].message.content.strip() except Exception as e: - return f"Error generating response: {e!s}" + return f"Error generating response: {str(e)}" @log -def main() -> None: +def main(): print("Out-of-Context RAG Demo") print("This demo shows how RAG systems can generate out-of-context information and how to prevent it.") @@ -197,7 +210,7 @@ def main() -> None: print("\nSuggested queries (these will demonstrate the problem):") for i, q in enumerate(suggested_queries): - print(f"{i + 1}. {q}") + print(f"{i+1}. {q}") # Main interaction loop while True: @@ -234,7 +247,7 @@ def main() -> None: break except Exception as e: - print(f"Error: {e!s}") + print(f"Error: {str(e)}") if __name__ == "__main__": diff --git a/examples/rag/cli-rag-demo/test.py b/examples/rag/cli-rag-demo/test.py index c5a96ff6..8cc1d94f 100644 --- a/examples/rag/cli-rag-demo/test.py +++ b/examples/rag/cli-rag-demo/test.py @@ -1,8 +1,7 @@ import os - from dotenv import load_dotenv -from splunk_ao import openai, splunk_ao_context +from splunk_ao import splunk_ao_context, openai load_dotenv() @@ -14,9 +13,7 @@ def call_openai(): - chat_completion = client.chat.completions.create( - messages=[{"role": "user", "content": "Say this is a test"}], model="gpt-4o" - ) + chat_completion = client.chat.completions.create(messages=[{"role": "user", "content": "Say this is a test"}], model="gpt-4o") return chat_completion.choices[0].message.content diff --git a/examples/rag/elastic-chatbot-rag-app/api/app.py b/examples/rag/elastic-chatbot-rag-app/api/app.py index e52ecbf3..a94f4a24 100644 --- a/examples/rag/elastic-chatbot-rag-app/api/app.py +++ b/examples/rag/elastic-chatbot-rag-app/api/app.py @@ -27,7 +27,7 @@ def api_chat(): @app.cli.command() -def create_index() -> None: +def create_index(): """Create or re-create the Elasticsearch index.""" basedir = os.path.abspath(os.path.dirname(__file__)) sys.path.append(f"{basedir}/../") diff --git a/examples/rag/elastic-chatbot-rag-app/api/chat.py b/examples/rag/elastic-chatbot-rag-app/api/chat.py index eef22891..7d790f6f 100644 --- a/examples/rag/elastic-chatbot-rag-app/api/chat.py +++ b/examples/rag/elastic-chatbot-rag-app/api/chat.py @@ -1,16 +1,25 @@ import json import os -from functools import cache -from elasticsearch_client import elasticsearch_client, get_elasticsearch_chat_message_history +from elasticsearch_client import ( + elasticsearch_client, + get_elasticsearch_chat_message_history, +) from flask import current_app, render_template, stream_with_context +from functools import cache +from langchain_elasticsearch import ( + ElasticsearchStore, + SparseVectorStrategy, +) from langchain_core.messages.ai import AIMessageChunk -from langchain_elasticsearch import ElasticsearchStore, SparseVectorStrategy + from langgraph.graph import START, StateGraph -from llm_integrations import get_llm +from splunk_ao.handlers.langchain import SplunkAOCallback + + from utils import State -from splunk_ao.handlers.langchain import SplunkAOCallback +from llm_integrations import get_llm # Make sure to set your Splunk AO logging env variables callback = SplunkAOCallback() @@ -24,7 +33,9 @@ DONE_TAG = "[DONE]" store = ElasticsearchStore( - es_connection=elasticsearch_client, index_name=INDEX, strategy=SparseVectorStrategy(model_id=ELSER_MODEL) + es_connection=elasticsearch_client, + index_name=INDEX, + strategy=SparseVectorStrategy(model_id=ELSER_MODEL), ) store_retriever = store.as_retriever() @@ -50,7 +61,9 @@ def retrieve(state: State): if len(chat_history.messages) > 0: # create a condensed question condense_question_prompt = render_template( - "condense_question_prompt.txt", question=question, chat_history=chat_history.messages + "condense_question_prompt.txt", + question=question, + chat_history=chat_history.messages, ) condensed_question = llm.invoke(condense_question_prompt).content else: @@ -66,7 +79,12 @@ def generate(state: State): question = state["question"] docs = state["context"] - qa_prompt = render_template("rag_prompt.txt", question=question, docs=docs, chat_history=chat_history.messages) + qa_prompt = render_template( + "rag_prompt.txt", + question=question, + docs=docs, + chat_history=chat_history.messages, + ) response = llm.invoke(qa_prompt) return {"answer": response.content} @@ -77,7 +95,9 @@ def generate(state: State): retrieved = False answer = "" for mode, step in graph.stream( - {"question": question}, stream_mode=["updates", "messages"], config={"callbacks": [callback]} + {"question": question}, + stream_mode=["updates", "messages"], + config={"callbacks": [callback]}, ): if mode == "updates": if not retrieved and "retrieve" in step: diff --git a/examples/rag/elastic-chatbot-rag-app/api/elasticsearch_client.py b/examples/rag/elastic-chatbot-rag-app/api/elasticsearch_client.py index f5ed47bf..a76aa5e0 100644 --- a/examples/rag/elastic-chatbot-rag-app/api/elasticsearch_client.py +++ b/examples/rag/elastic-chatbot-rag-app/api/elasticsearch_client.py @@ -10,7 +10,8 @@ if ELASTICSEARCH_USER: elasticsearch_client = Elasticsearch( - hosts=[ELASTICSEARCH_URL], basic_auth=(ELASTICSEARCH_USER, ELASTICSEARCH_PASSWORD) + hosts=[ELASTICSEARCH_URL], + basic_auth=(ELASTICSEARCH_USER, ELASTICSEARCH_PASSWORD), ) elif ELASTICSEARCH_API_KEY: elasticsearch_client = Elasticsearch(hosts=[ELASTICSEARCH_URL], api_key=ELASTICSEARCH_API_KEY) diff --git a/examples/rag/elastic-chatbot-rag-app/api/llm_integrations.py b/examples/rag/elastic-chatbot-rag-app/api/llm_integrations.py index fdd97a4f..5b371ee1 100644 --- a/examples/rag/elastic-chatbot-rag-app/api/llm_integrations.py +++ b/examples/rag/elastic-chatbot-rag-app/api/llm_integrations.py @@ -1,5 +1,4 @@ import os - from langchain.chat_models import init_chat_model LLM_TYPE = os.getenv("LLM_TYPE", "openai") @@ -60,8 +59,6 @@ def init_anthropic_chat(temperature): def get_llm(temperature=0): if LLM_TYPE not in MAP_LLM_TYPE_TO_CHAT_MODEL: - raise Exception( - "LLM type not found. Please set LLM_TYPE to one of: " + ", ".join(MAP_LLM_TYPE_TO_CHAT_MODEL.keys()) + "." - ) + raise Exception("LLM type not found. Please set LLM_TYPE to one of: " + ", ".join(MAP_LLM_TYPE_TO_CHAT_MODEL.keys()) + ".") return MAP_LLM_TYPE_TO_CHAT_MODEL[LLM_TYPE](temperature=temperature) diff --git a/examples/rag/elastic-chatbot-rag-app/api/utils.py b/examples/rag/elastic-chatbot-rag-app/api/utils.py index a5f45ecf..25b5ee3f 100644 --- a/examples/rag/elastic-chatbot-rag-app/api/utils.py +++ b/examples/rag/elastic-chatbot-rag-app/api/utils.py @@ -1,9 +1,9 @@ +from typing_extensions import List, TypedDict from langchain_core.documents import Document -from typing_extensions import TypedDict # Define state for application class State(TypedDict): question: str - context: list[Document] + context: List[Document] answer: str diff --git a/examples/rag/elastic-chatbot-rag-app/data/index_data.py b/examples/rag/elastic-chatbot-rag-app/data/index_data.py index 808d0508..0088da2d 100644 --- a/examples/rag/elastic-chatbot-rag-app/data/index_data.py +++ b/examples/rag/elastic-chatbot-rag-app/data/index_data.py @@ -1,12 +1,18 @@ import json import os -import time from sys import stdout +import time +from halo import Halo from warnings import warn +from elasticsearch import ( + ApiError, + Elasticsearch, + NotFoundError, + BadRequestError, +) from elastic_transport._exceptions import ConnectionTimeout -from elasticsearch import ApiError, BadRequestError, Elasticsearch, NotFoundError -from halo import Halo + from langchain.docstore.document import Document from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_elasticsearch import ElasticsearchStore @@ -22,14 +28,17 @@ ELSER_MODEL = os.getenv("ELSER_MODEL", ".elser_model_2") if ELASTICSEARCH_USER: - es = Elasticsearch(hosts=[ELASTICSEARCH_URL], basic_auth=(ELASTICSEARCH_USER, ELASTICSEARCH_PASSWORD)) + es = Elasticsearch( + hosts=[ELASTICSEARCH_URL], + basic_auth=(ELASTICSEARCH_USER, ELASTICSEARCH_PASSWORD), + ) elif ELASTICSEARCH_API_KEY: es = Elasticsearch(hosts=[ELASTICSEARCH_URL], api_key=ELASTICSEARCH_API_KEY) else: raise ValueError("Please provide either ELASTICSEARCH_USER or ELASTICSEARCH_API_KEY") -def install_elser() -> None: +def install_elser(): # This script is re-entered on ctrl-c or someone just running it twice. # Hence, both steps need to be careful about being potentially redundant. @@ -65,17 +74,20 @@ def is_elser_fully_allocated(): return allocation_status.get("state") == "fully_allocated" -def main() -> None: +def main(): install_elser() print(f"Loading data from ${FILE}") metadata_keys = ["name", "summary", "url", "category", "updated_at"] workplace_docs = [] - with open(FILE) as f: + with open(FILE, "rt") as f: for doc in json.loads(f.read()): workplace_docs.append( - Document(page_content=doc["content"], metadata={k: doc.get(k) for k in metadata_keys}) + Document( + page_content=doc["content"], + metadata={k: doc.get(k) for k in metadata_keys}, + ) ) print(f"Loaded {len(workplace_docs)} documents") @@ -116,7 +128,7 @@ def main() -> None: except (ConnectionTimeout, ApiError) as e: if isinstance(e, ApiError) and e.status_code != 408: raise - warn(f"Error occurred, will retry after ML jobs complete: {e}", stacklevel=2) + warn(f"Error occurred, will retry after ML jobs complete: {e}") await_ml_tasks() es.indices.delete(index=INDEX, ignore_unavailable=True) store.add_documents(list(docs)) @@ -127,7 +139,7 @@ def main() -> None: print(f"Documents added to index {INDEX}") -def await_ml_tasks(max_timeout=1200, interval=5) -> None: +def await_ml_tasks(max_timeout=1200, interval=5): """ Waits for all machine learning tasks to complete within a specified timeout period. diff --git a/poetry.lock b/poetry.lock index 3fcb5eeb..d9e511f8 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -173,19 +173,6 @@ files = [ frozenlist = ">=1.1.0" typing-extensions = {version = ">=4.2", markers = "python_version < \"3.13\""} -[[package]] -name = "annotated-doc" -version = "0.0.4" -description = "Document parameters, class attributes, return types, and variables inline, with Annotated." -optional = false -python-versions = ">=3.8" -groups = ["main", "dev"] -files = [ - {file = "annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320"}, - {file = "annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4"}, -] -markers = {main = "(extra == \"crewai\" or extra == \"all\") and python_version <= \"3.13\""} - [[package]] name = "annotated-types" version = "0.7.0" @@ -467,7 +454,7 @@ files = [ {file = "cffi-1.17.1-cp39-cp39-win_amd64.whl", hash = "sha256:d016c76bdd850f3c626af19b0542c9677ba156e4ee4fccfdd7848803533ef662"}, {file = "cffi-1.17.1.tar.gz", hash = "sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824"}, ] -markers = {main = "platform_python_implementation != \"PyPy\" and (extra == \"openai\" or extra == \"all\") or platform_python_implementation != \"PyPy\" and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") or (extra == \"langchain\" or extra == \"all\") and platform_python_implementation == \"PyPy\" or (extra == \"langchain\" or extra == \"all\") and (extra == \"openai\" or extra == \"all\") or (extra == \"langchain\" or extra == \"all\") and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation == \"PyPy\""} +markers = {main = "extra == \"openai\" or extra == \"all\" or platform_python_implementation == \"PyPy\" or (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") and python_version <= \"3.13\"", test = "platform_python_implementation == \"PyPy\""} [package.dependencies] pycparser = "*" @@ -572,7 +559,6 @@ files = [ {file = "charset_normalizer-3.4.3-py3-none-any.whl", hash = "sha256:ce571ab16d890d23b5c278547ba694193a45011ff86a9162a71307ed9f86759a"}, {file = "charset_normalizer-3.4.3.tar.gz", hash = "sha256:6fce4b8500244f6fcb71465d4a4930d132ba9ab8e71a7859e6a5d59851068d14"}, ] -markers = {main = "python_version < \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or extra == \"crewai\")"} [[package]] name = "chromadb" @@ -627,14 +613,14 @@ dev = ["chroma-hnswlib (==0.7.6)", "fastapi (>=0.115.9)", "opentelemetry-instrum name = "click" version = "8.1.8" description = "Composable command line interface toolkit" -optional = true +optional = false python-versions = ">=3.7" -groups = ["main"] -markers = "sys_platform != \"emscripten\" and (extra == \"openai\" or extra == \"all\") or python_version <= \"3.13\" and sys_platform != \"emscripten\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +groups = ["main", "dev"] files = [ {file = "click-8.1.8-py3-none-any.whl", hash = "sha256:63c132bbbed01578a06712a2d1f497bb62d9c1c0d329b7903a866228027263b2"}, {file = "click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a"}, ] +markers = {main = "sys_platform != \"emscripten\" and (extra == \"openai\" or extra == \"all\") or python_version <= \"3.13\" and sys_platform != \"emscripten\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")"} [package.dependencies] colorama = {version = "*", markers = "platform_system == \"Windows\""} @@ -914,7 +900,6 @@ files = [ {file = "distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2"}, {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, ] -markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"openai\") and python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\""} [[package]] name = "docstring-parser" @@ -923,11 +908,11 @@ description = "Parse Python docstrings in reST, Google and Numpydoc format" optional = false python-versions = ">=3.8" groups = ["main", "test"] +markers = {main = "(extra == \"crewai\" or extra == \"all\") and python_version <= \"3.13\""} files = [ {file = "docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708"}, {file = "docstring_parser-0.17.0.tar.gz", hash = "sha256:583de4a309722b3315439bb31d64ba3eebada841f2e2cee23b99df001434c912"}, ] -markers = {main = "(extra == \"crewai\" or extra == \"all\") and python_version <= \"3.13\""} [package.extras] dev = ["pre-commit (>=2.16.0) ; python_version >= \"3.9\"", "pydoctor (>=25.4.0)", "pytest"] @@ -1891,7 +1876,6 @@ files = [ {file = "jiter-0.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:1b28302349dc65703a9e4ead16f163b1c339efffbe1049c30a44b001a2a4fff9"}, {file = "jiter-0.10.0.tar.gz", hash = "sha256:07a7142c38aacc85194391108dc91b5b57093c978a9932bd86a36862759d9500"}, ] -markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"openai\") and python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\""} [[package]] name = "json-repair" @@ -1933,7 +1917,6 @@ files = [ {file = "jsonpatch-1.33-py2.py3-none-any.whl", hash = "sha256:0ae28c0cd062bbd8b8ecc26d7d164fbbea9652a1a3693f3b956c1eae5145dade"}, {file = "jsonpatch-1.33.tar.gz", hash = "sha256:9fcd4009c41e6d12348b4a0ff2563ba56a2923a7dfee731d004e212e1ee5030c"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] jsonpointer = ">=1.9" @@ -1949,7 +1932,6 @@ files = [ {file = "jsonpointer-3.0.0-py2.py3-none-any.whl", hash = "sha256:13e088adc14fca8b6aa8177c044e12701e6ad4b28ff10e65f2267a90109c9942"}, {file = "jsonpointer-3.0.0.tar.gz", hash = "sha256:2b2d729f2091522d61c3b31f82e11870f60b68f43fbc705cb76bf4b832af59ef"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\""} [[package]] name = "jsonref" @@ -1979,7 +1961,7 @@ files = [ [package.dependencies] attrs = ">=22.2.0" -jsonschema-specifications = ">=2023.3.6" +jsonschema-specifications = ">=2023.03.6" referencing = ">=0.28.4" rpds-py = ">=0.7.1" @@ -2017,7 +1999,7 @@ files = [ ] [package.dependencies] -certifi = ">=14.5.14" +certifi = ">=14.05.14" durationpy = ">=0.7" google-auth = ">=1.0.1" oauthlib = ">=3.2.2" @@ -2079,7 +2061,6 @@ files = [ {file = "langchain_core-1.2.7-py3-none-any.whl", hash = "sha256:452f4fef7a3d883357b22600788d37e3d8854ef29da345b7ac7099f33c31828b"}, {file = "langchain_core-1.2.7.tar.gz", hash = "sha256:e1460639f96c352b4a41c375f25aeb8d16ffc1769499fb1c20503aad59305ced"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] jsonpatch = ">=1.33.0,<2.0.0" @@ -2174,7 +2155,6 @@ files = [ {file = "langsmith-0.4.14-py3-none-any.whl", hash = "sha256:b6d070ac425196947d2a98126fb0e35f3b8c001a2e6e5b7049dd1c56f0767d0b"}, {file = "langsmith-0.4.14.tar.gz", hash = "sha256:4d29c7a9c85b20ba813ab9c855407bccdf5eb4f397f512ffa89959b2a2cb83ed"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] httpx = ">=0.23.0,<1" @@ -2872,7 +2852,6 @@ files = [ {file = "openai-2.32.0-py3-none-any.whl", hash = "sha256:4dcc9badeb4bf54ad0d187453742f290226d30150890b7890711bda4f32f192f"}, {file = "openai-2.32.0.tar.gz", hash = "sha256:c54b27a9e4cb8d51f0dd94972ffd1a04437efeb259a9e60d8922b8bd26fe55e0"}, ] -markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"openai\") and python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\""} [package.dependencies] anyio = ">=3.5.0,<5" @@ -2937,26 +2916,28 @@ voice = ["numpy (>=2.2.0,<3) ; python_version >= \"3.10\"", "websockets (>=15.0, [[package]] name = "openapi-python-client" -version = "0.29.0" +version = "0.26.1" description = "Generate modern Python clients from OpenAPI" optional = false -python-versions = "<4.0,>=3.11" +python-versions = "<4.0,>=3.9" groups = ["dev"] files = [ - {file = "openapi_python_client-0.29.0-py3-none-any.whl", hash = "sha256:1085864c1e0a42fb50e1f5eb84363b19de07ebfb8a3f82a146c8529b948b8f12"}, - {file = "openapi_python_client-0.29.0.tar.gz", hash = "sha256:4ff439a63765aaef548e69c4eedd86dfad57891dc322709450af0c2c9a72a23e"}, + {file = "openapi_python_client-0.26.1-py3-none-any.whl", hash = "sha256:415cb8095b1a3f15cec45670c5075c5097f65390a351d21512e8f6ea5c1be644"}, + {file = "openapi_python_client-0.26.1.tar.gz", hash = "sha256:e3832f0ef074a0ab591d1eeb5d3dab2ca820cd0349f7e79d9663b7b21206be5d"}, ] [package.dependencies] attrs = ">=22.2.0" colorama = {version = ">=0.4.3", markers = "sys_platform == \"win32\""} -httpx = ">=0.23.1,<0.29.0" +httpx = ">=0.23.0,<0.29.0" jinja2 = ">=3.0.0,<4.0.0" pydantic = ">=2.10,<3.0.0" -ruamel-yaml = ">=0.18.6,<0.20.0" -ruff = ">=0.2" +python-dateutil = ">=2.8.1,<3.0.0" +ruamel-yaml = ">=0.18.6,<0.19.0" +ruff = ">=0.2,<0.14" shellingham = ">=1.3.2,<2.0.0" -typer = ">0.16,<0.27" +typer = ">0.6,<0.18" +typing-extensions = ">=4.8.0,<5.0.0" [[package]] name = "openpyxl" @@ -3214,7 +3195,7 @@ files = [ {file = "orjson-3.11.2-cp39-cp39-win_amd64.whl", hash = "sha256:c9ec0cc0d4308cad1e38a1ee23b64567e2ff364c2a3fe3d6cbc69cf911c45712"}, {file = "orjson-3.11.2.tar.gz", hash = "sha256:91bdcf5e69a8fd8e8bdb3de32b31ff01d2bd60c1e8d5fe7d5afabdcf19920309"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") and python_version <= \"3.13\"", test = "platform_python_implementation != \"PyPy\""} +markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"langchain\" or extra == \"all\" or (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\") and python_version <= \"3.13\"", test = "platform_python_implementation != \"PyPy\""} [[package]] name = "ormsgpack" @@ -3299,7 +3280,6 @@ files = [ {file = "packaging-24.2-py3-none-any.whl", hash = "sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759"}, {file = "packaging-24.2.tar.gz", hash = "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f"}, ] -markers = {main = "python_version < \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"crewai\")"} [[package]] name = "pathspec" @@ -3965,7 +3945,7 @@ files = [ {file = "pycparser-2.22-py3-none-any.whl", hash = "sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc"}, {file = "pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6"}, ] -markers = {main = "platform_python_implementation != \"PyPy\" and (extra == \"openai\" or extra == \"all\") or platform_python_implementation != \"PyPy\" and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") or (extra == \"langchain\" or extra == \"all\") and platform_python_implementation == \"PyPy\" or (extra == \"langchain\" or extra == \"all\") and (extra == \"openai\" or extra == \"all\") or (extra == \"langchain\" or extra == \"all\") and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation == \"PyPy\""} +markers = {main = "extra == \"openai\" or extra == \"all\" or platform_python_implementation == \"PyPy\" or (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") and python_version <= \"3.13\"", test = "platform_python_implementation == \"PyPy\""} [[package]] name = "pydantic" @@ -4420,7 +4400,7 @@ version = "2.9.0.post0" description = "Extensions to the standard Python datetime module" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -groups = ["main", "test"] +groups = ["main", "dev", "test"] files = [ {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"}, {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"}, @@ -4550,7 +4530,6 @@ files = [ {file = "PyYAML-6.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8"}, {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"}, ] -markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") and python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\""} [[package]] name = "referencing" @@ -4679,7 +4658,6 @@ files = [ {file = "requests-2.32.4-py3-none-any.whl", hash = "sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c"}, {file = "requests-2.32.4.tar.gz", hash = "sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422"}, ] -markers = {main = "python_version < \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or extra == \"crewai\")"} [package.dependencies] certifi = ">=2017.4.17" @@ -4740,7 +4718,6 @@ files = [ {file = "requests-toolbelt-1.0.0.tar.gz", hash = "sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6"}, {file = "requests_toolbelt-1.0.0-py2.py3-none-any.whl", hash = "sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] requests = ">=2.0.1,<3.0.0" @@ -5040,30 +5017,30 @@ files = [ [[package]] name = "ruff" -version = "0.15.20" +version = "0.12.8" description = "An extremely fast Python linter and code formatter, written in Rust." optional = false python-versions = ">=3.7" groups = ["dev"] files = [ - {file = "ruff-0.15.20-py3-none-linux_armv6l.whl", hash = "sha256:00e188c53e499c3c1637f73c91dcf2fb56d576cab76ce1be50a27c4e80e37078"}, - {file = "ruff-0.15.20-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9ebd1fd9b9c95fc0bd7b2761aebec1f030013d2e193a2901b224af68fe47251b"}, - {file = "ruff-0.15.20-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c5b16cdd67ca108185cd36dce98c576350c03b1660a751de725fb049193a0632"}, - {file = "ruff-0.15.20-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3413bb3c3d2ca6a8208f1f4809cd2dca3c6de6d0b491c0e70847672bde6e6efd"}, - {file = "ruff-0.15.20-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bd7ec42b3bb3da066488db093308a69c4ac5ee6d2af333a86ba6e2eb2e7dd44b"}, - {file = "ruff-0.15.20-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e1a36ad0eb77fba9aabfb69ede54de6f376d04ac18ebea022847046d340a8267"}, - {file = "ruff-0.15.20-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b6df3b1e4610432f0386dba04d853b5f08cbbc903410c6fcc02f620f05aff53c"}, - {file = "ruff-0.15.20-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e89f198a1ea6ef0d727c1cf16088bc91a6cb0ab947dedc966715691647186eae"}, - {file = "ruff-0.15.20-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:309809086c2acb67624950a3c8133e80f32d0d3e27106c0cd60ff26657c9f24b"}, - {file = "ruff-0.15.20-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:2d2374caa2f2c2f9e2b7da0a50802cfb8b79f55a9b5e49379f564544fbf56487"}, - {file = "ruff-0.15.20-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:a1ed17b65293e0c2f22fc387bc13198a5de94bf4429589b0ff6946b0feaf21a3"}, - {file = "ruff-0.15.20-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:f701305e66b38ea6c91882490eb73459796808e4c6362a1b765255e0cdcd4053"}, - {file = "ruff-0.15.20-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5b9c0c367ad8e5d0d5b5b8537864c469a0a0e55417aadfbeca41fa61333be9f4"}, - {file = "ruff-0.15.20-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:01cc00dd58f0df339d0e902219dd53990ea99996a0344e5d9cc8d45d5307e460"}, - {file = "ruff-0.15.20-py3-none-win32.whl", hash = "sha256:ed65ef510e43a137207e0f01cfcf998aeddb1aeeda5c9d35023e910284d7cf21"}, - {file = "ruff-0.15.20-py3-none-win_amd64.whl", hash = "sha256:a525c81c70fb0380344dd1d8745d8cc1c890b7fc94a58d5a07bd8eb9557b8415"}, - {file = "ruff-0.15.20-py3-none-win_arm64.whl", hash = "sha256:2f5b2a6d614e8700388806a14996c40fab2c47b819ef57d790a34878858ed9ca"}, - {file = "ruff-0.15.20.tar.gz", hash = "sha256:1416eb04349192646b54de98f146c4f59afe37d0decfc02c3cbbf396f3a28566"}, + {file = "ruff-0.12.8-py3-none-linux_armv6l.whl", hash = "sha256:63cb5a5e933fc913e5823a0dfdc3c99add73f52d139d6cd5cc8639d0e0465513"}, + {file = "ruff-0.12.8-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9a9bbe28f9f551accf84a24c366c1aa8774d6748438b47174f8e8565ab9dedbc"}, + {file = "ruff-0.12.8-py3-none-macosx_11_0_arm64.whl", hash = "sha256:2fae54e752a3150f7ee0e09bce2e133caf10ce9d971510a9b925392dc98d2fec"}, + {file = "ruff-0.12.8-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c0acbcf01206df963d9331b5838fb31f3b44fa979ee7fa368b9b9057d89f4a53"}, + {file = "ruff-0.12.8-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ae3e7504666ad4c62f9ac8eedb52a93f9ebdeb34742b8b71cd3cccd24912719f"}, + {file = "ruff-0.12.8-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cb82efb5d35d07497813a1c5647867390a7d83304562607f3579602fa3d7d46f"}, + {file = "ruff-0.12.8-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:dbea798fc0065ad0b84a2947b0aff4233f0cb30f226f00a2c5850ca4393de609"}, + {file = "ruff-0.12.8-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:49ebcaccc2bdad86fd51b7864e3d808aad404aab8df33d469b6e65584656263a"}, + {file = "ruff-0.12.8-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0ac9c570634b98c71c88cb17badd90f13fc076a472ba6ef1d113d8ed3df109fb"}, + {file = "ruff-0.12.8-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:560e0cd641e45591a3e42cb50ef61ce07162b9c233786663fdce2d8557d99818"}, + {file = "ruff-0.12.8-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:71c83121512e7743fba5a8848c261dcc454cafb3ef2934a43f1b7a4eb5a447ea"}, + {file = "ruff-0.12.8-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:de4429ef2ba091ecddedd300f4c3f24bca875d3d8b23340728c3cb0da81072c3"}, + {file = "ruff-0.12.8-py3-none-musllinux_1_2_i686.whl", hash = "sha256:a2cab5f60d5b65b50fba39a8950c8746df1627d54ba1197f970763917184b161"}, + {file = "ruff-0.12.8-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:45c32487e14f60b88aad6be9fd5da5093dbefb0e3e1224131cb1d441d7cb7d46"}, + {file = "ruff-0.12.8-py3-none-win32.whl", hash = "sha256:daf3475060a617fd5bc80638aeaf2f5937f10af3ec44464e280a9d2218e720d3"}, + {file = "ruff-0.12.8-py3-none-win_amd64.whl", hash = "sha256:7209531f1a1fcfbe8e46bcd7ab30e2f43604d8ba1c49029bb420b103d0b5f76e"}, + {file = "ruff-0.12.8-py3-none-win_arm64.whl", hash = "sha256:c90e1a334683ce41b0e7a04f41790c429bf5073b62c1ae701c9dc5b3d14f0749"}, + {file = "ruff-0.12.8.tar.gz", hash = "sha256:4cb3a45525176e1009b2b64126acf5f9444ea59066262791febf55e40493a033"}, ] [[package]] @@ -5085,7 +5062,7 @@ version = "1.17.0" description = "Python 2 and 3 compatibility utilities" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -groups = ["main", "test"] +groups = ["main", "dev", "test"] files = [ {file = "six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274"}, {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, @@ -5174,7 +5151,6 @@ files = [ {file = "tenacity-9.1.2-py3-none-any.whl", hash = "sha256:f77bf36710d8b73a50b2dd155c97b870017ad21afe6ab300326b0371b3b05138"}, {file = "tenacity-9.1.2.tar.gz", hash = "sha256:1169d376c297e7de388d18b4481760d478b0e99a777cad3a9c86e556f4b697cb"}, ] -markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") and python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\""} [package.extras] doc = ["reno", "sphinx"] @@ -5475,22 +5451,22 @@ telegram = ["requests"] [[package]] name = "typer" -version = "0.26.8" +version = "0.16.0" description = "Typer, build great CLIs. Easy to code. Based on Python type hints." optional = false -python-versions = ">=3.10" +python-versions = ">=3.7" groups = ["main", "dev"] files = [ - {file = "typer-0.26.8-py3-none-any.whl", hash = "sha256:3512ca79ac5c11113414b36e80281b872884477722440691c89d1112e321a49c"}, - {file = "typer-0.26.8.tar.gz", hash = "sha256:c244a6bd558886fe3f8780efb6bdd28bb9aff005a94eedebaa5cb32926fe2f7e"}, + {file = "typer-0.16.0-py3-none-any.whl", hash = "sha256:1f79bed11d4d02d4310e3c1b7ba594183bcedb0ac73b27a9e5f28f6fb5b98855"}, + {file = "typer-0.16.0.tar.gz", hash = "sha256:af377ffaee1dbe37ae9440cb4e8f11686ea5ce4e9bae01b84ae7c63b87f1dd3b"}, ] markers = {main = "(extra == \"crewai\" or extra == \"all\") and python_version <= \"3.13\""} [package.dependencies] -annotated-doc = ">=0.0.2" -colorama = {version = "*", markers = "platform_system == \"Windows\""} -rich = ">=13.8.0" +click = ">=8.0.0" +rich = ">=10.11.0" shellingham = ">=1.3.0" +typing-extensions = ">=3.7.4.3" [[package]] name = "types-requests" @@ -5571,11 +5547,11 @@ description = "HTTP library with thread-safe connection pooling, file post, and optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7" groups = ["main", "test"] +markers = "platform_python_implementation == \"PyPy\"" files = [ {file = "urllib3-1.26.20-py2.py3-none-any.whl", hash = "sha256:0ed14ccfbf1c30a9072c7ca157e4319b70d65f623e91e7b32fadb2853431016e"}, {file = "urllib3-1.26.20.tar.gz", hash = "sha256:40c2dc0c681e47eb8f90e7e27bf6ff7df2e677421fd46756da1161c39ca70d32"}, ] -markers = {main = "platform_python_implementation == \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\") or platform_python_implementation == \"PyPy\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") and python_version <= \"3.13\"", test = "platform_python_implementation == \"PyPy\""} [package.extras] brotli = ["brotli (==1.0.9) ; os_name != \"nt\" and python_version < \"3\" and platform_python_implementation == \"CPython\"", "brotli (>=1.0.9) ; python_version >= \"3\" and platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; (os_name != \"nt\" or python_version >= \"3\") and platform_python_implementation != \"CPython\"", "brotlipy (>=0.6.0) ; os_name == \"nt\" and python_version < \"3\""] @@ -5589,11 +5565,11 @@ description = "HTTP library with thread-safe connection pooling, file post, and optional = false python-versions = ">=3.9" groups = ["main", "test"] +markers = "platform_python_implementation != \"PyPy\"" files = [ {file = "urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc"}, {file = "urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760"}, ] -markers = {main = "platform_python_implementation != \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\") or platform_python_implementation != \"PyPy\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") and python_version <= \"3.13\"", test = "platform_python_implementation != \"PyPy\""} [package.extras] brotli = ["brotli (>=1.0.9) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\""] @@ -5632,7 +5608,6 @@ files = [ {file = "uuid_utils-0.13.0-pp311-pypy311_pp73-manylinux_2_24_x86_64.whl", hash = "sha256:b7ccaa20e24c5f60f41a69ef571ed820737f9b0ade4cbeef56aaa8f80f5aa475"}, {file = "uuid_utils-0.13.0.tar.gz", hash = "sha256:4c17df6427a9e23a4cd7fb9ee1efb53b8abb078660b9bdb2524ca8595022dfe1"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\""} [[package]] name = "uv" @@ -5759,8 +5734,8 @@ files = [ [package.dependencies] PyYAML = "*" urllib3 = [ - {version = "*", markers = "platform_python_implementation != \"PyPy\" and python_version >= \"3.10\""}, {version = "<2", markers = "platform_python_implementation == \"PyPy\""}, + {version = "*", markers = "platform_python_implementation != \"PyPy\" and python_version >= \"3.10\""}, ] wrapt = "*" yarl = "*" @@ -6529,7 +6504,6 @@ files = [ {file = "zstandard-0.23.0-cp39-cp39-win_amd64.whl", hash = "sha256:f8346bfa098532bc1fb6c7ef06783e969d87a99dd1d2a5a18a892c1d7a643c58"}, {file = "zstandard-0.23.0.tar.gz", hash = "sha256:b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] cffi = {version = ">=1.11", markers = "platform_python_implementation == \"PyPy\""} diff --git a/pyproject.toml b/pyproject.toml index fb68b278..75dcb7c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,8 +79,8 @@ mypy = "^1.16.0" invoke = "^2.2.0" openai = ">=2.8.0,<3.0.0" fastapi = "^0.115.0" -ruff = "^0.15.20" -openapi-python-client = "^0.29.0" +ruff = "^0.12.3" +openapi-python-client = "^0.26.1" yq = "^3.4.3" pyyaml = "^6.0.2" pytest-timeout = "^2.4.0" @@ -260,26 +260,11 @@ ignore = [ "PT", # Pytest style rules (can be overly strict) "T201", # Print statements (allowed in tests for debugging) "D", # All docstring rules (not critical for tests) - "PLC0415", # Import not at top-level (needed for pytest.importorskip conditional imports) - "RUF043", # Regex metacharacters in match= (common in test parametrize patterns) - "E402", # Module level import not at top (needed after pytest.importorskip) ] # Allow print statements in examples (they're documentation) and scripts (utility tools) "examples/**/*.py" = [ "T201", # Print statements (allowed in examples) "D", # All docstring rules (examples focus on usage, not documentation) - "ANN", # Type annotations (examples are concise demos, not typed production code) - "PLC0415", # Import not at top-level (conditional/optional imports common in examples) - "E402", # Module level import not at top (common in example scripts) - "RUF003", # Ambiguous unicode character in comment (ℹ info symbol used intentionally) - "ASYNC", # Async rules (examples may use simplified async patterns) - "S", # Security rules (examples are not production code) - "RUF012", # Mutable class attributes (examples use simple patterns) - "ISC001", # Implicit string concatenation (style preference in examples) - "E722", # Bare except (examples often use simplified error handling) - "SIM", # Simplify rules (examples prioritize clarity over brevity) - "B", # Bugbear rules (examples may use simplified patterns) - "ARG", # Unused arguments (examples often have placeholder callbacks) ] "scripts/**/*.py" = [ "T201", # Print statements (allowed in scripts) @@ -293,8 +278,6 @@ ignore = [ "src/splunk_ao/handlers/crewai/handler.py" = ["PLC0415"] # Version-specific conditional imports "src/splunk_ao/project.py" = ["PLC0415"] # Bottom-of-file circular import avoidance "src/splunk_ao/logger/logger.py" = ["PLC0415"] # Local imports to avoid circular dependencies -"src/splunk_ao/experiment.py" = ["PLC0415"] # Local imports to avoid circular dependencies -"src/splunk_ao/metric.py" = ["PLC0415"] # Local imports to avoid circular dependencies "src/splunk_ao/logger/__init__.py" = ["PLC0415"] # Lazy import for GalileoLogger [tool.ruff.lint.isort] diff --git a/scripts/create_docs.py b/scripts/create_docs.py index e9acad15..4a2fc7b1 100644 --- a/scripts/create_docs.py +++ b/scripts/create_docs.py @@ -553,7 +553,8 @@ def _convert_rst_code_blocks(text: str) -> str: def _sanitize_description(text: str) -> str: """Convert RST code blocks and escape MDX-unsafe curly braces in description text.""" text = _convert_rst_code_blocks(text) - return _escape_curly_braces(text) + text = _escape_curly_braces(text) + return text def write_function(parts: list[str], fn: Any, heading_level: int = 3) -> None: @@ -631,7 +632,7 @@ def write_function(parts: list[str], fn: Any, heading_level: int = 3) -> None: continue if "```" in code: parts.append(code) - elif code.startswith((">>>", "...")): + elif code.startswith(">>>") or code.startswith("..."): parts.append("```python") parts.append(code[4:]) parts.append("```") @@ -709,7 +710,7 @@ def write_class(parts: list[str], cls: Any) -> None: continue if "```" in code: parts.append(code) - elif code.startswith((">>>", "...")): + elif code.startswith(">>>") or code.startswith("..."): parts.append("```python") parts.append(code[4:]) parts.append("```") diff --git a/splunk-ao-a2a/examples/two_agent_demo.py b/splunk-ao-a2a/examples/two_agent_demo.py index 595d0166..3d5c5e2f 100644 --- a/splunk-ao-a2a/examples/two_agent_demo.py +++ b/splunk-ao-a2a/examples/two_agent_demo.py @@ -48,10 +48,10 @@ from langgraph.graph import END, START, StateGraph from opentelemetry.instrumentation.langchain import LangchainInstrumentor from opentelemetry.sdk.trace import TracerProvider +from splunk_ao.otel import SplunkAOSpanProcessor, add_splunk_ao_span_processor from starlette.applications import Starlette from typing_extensions import TypedDict -from splunk_ao.otel import SplunkAOSpanProcessor, add_splunk_ao_span_processor from splunk_ao_a2a import A2AInstrumentor load_dotenv(Path(__file__).parent / ".env") diff --git a/splunk-ao-a2a/src/splunk_ao_a2a/_context.py b/splunk-ao-a2a/src/splunk_ao_a2a/_context.py index e938aa47..12b9a1de 100644 --- a/splunk-ao-a2a/src/splunk_ao_a2a/_context.py +++ b/splunk-ao-a2a/src/splunk_ao_a2a/_context.py @@ -12,10 +12,10 @@ from splunk_ao_a2a._constants import ( AGNTCY_OBSERVE_KEY, + SPLUNK_AO_OBSERVE_KEY, LINK_FROM_AGENT, LINK_TYPE, LINK_TYPE_AGENT_HANDOFF, - SPLUNK_AO_OBSERVE_KEY, ) _logger = logging.getLogger(__name__) diff --git a/splunk-ao-adk/src/splunk_ao_adk/callback.py b/splunk-ao-adk/src/splunk_ao_adk/callback.py index a9858176..d973b5f9 100644 --- a/splunk-ao-adk/src/splunk_ao_adk/callback.py +++ b/splunk-ao-adk/src/splunk_ao_adk/callback.py @@ -8,6 +8,7 @@ from typing import Any from splunk_ao.schema.trace import TracesIngestRequest + from splunk_ao_adk.observer import ( SplunkAOObserver, get_agent_name_from_tool_context, diff --git a/splunk-ao-adk/src/splunk_ao_adk/observer.py b/splunk-ao-adk/src/splunk_ao_adk/observer.py index 759ce330..b21bd43f 100644 --- a/splunk-ao-adk/src/splunk_ao_adk/observer.py +++ b/splunk-ao-adk/src/splunk_ao_adk/observer.py @@ -14,6 +14,7 @@ from splunk_ao.handlers.base_handler import SplunkAOBaseHandler from splunk_ao.schema.trace import TracesIngestRequest from splunk_ao.utils.serialization import serialize_to_str + from splunk_ao_adk.data_converters import ( convert_adk_content_to_splunk_ao_messages, convert_adk_tools_to_splunk_ao_format, diff --git a/splunk-ao-adk/src/splunk_ao_adk/plugin.py b/splunk-ao-adk/src/splunk_ao_adk/plugin.py index da071f03..e2fa27d9 100644 --- a/splunk-ao-adk/src/splunk_ao_adk/plugin.py +++ b/splunk-ao-adk/src/splunk_ao_adk/plugin.py @@ -10,6 +10,7 @@ from uuid import UUID from splunk_ao.schema.trace import TracesIngestRequest + from splunk_ao_adk.observer import ( SplunkAOObserver, get_agent_name_from_tool_context, diff --git a/splunk-ao-adk/src/splunk_ao_adk/span_manager.py b/splunk-ao-adk/src/splunk_ao_adk/span_manager.py index 3af0df9a..b8ed09c9 100644 --- a/splunk-ao-adk/src/splunk_ao_adk/span_manager.py +++ b/splunk-ao-adk/src/splunk_ao_adk/span_manager.py @@ -8,6 +8,7 @@ from uuid import UUID from splunk_ao.handlers.base_handler import SplunkAOBaseHandler + from splunk_ao_adk.types import RunContext # Integration tag for all spans diff --git a/splunk-ao-adk/src/splunk_ao_adk/trace_builder.py b/splunk-ao-adk/src/splunk_ao_adk/trace_builder.py index 965b935e..7feb7d57 100644 --- a/splunk-ao-adk/src/splunk_ao_adk/trace_builder.py +++ b/splunk-ao-adk/src/splunk_ao_adk/trace_builder.py @@ -25,7 +25,6 @@ from galileo_core.schemas.logging.step import Metrics from galileo_core.schemas.shared.traces_logger import TracesLogger from pydantic import PrivateAttr - from splunk_ao.schema.logged import LoggedAgentSpan, LoggedLlmSpan, LoggedTrace, LoggedWorkflowSpan from splunk_ao.schema.trace import TracesIngestRequest from splunk_ao.utils.retrievers import convert_to_documents diff --git a/splunk-ao-adk/tests/conftest.py b/splunk-ao-adk/tests/conftest.py index 5c82b05f..231bd69a 100644 --- a/splunk-ao-adk/tests/conftest.py +++ b/splunk-ao-adk/tests/conftest.py @@ -8,10 +8,9 @@ from galileo_core.constants.routes import Routes as CoreRoutes from galileo_core.schemas.core.user import User from galileo_core.schemas.core.user_role import UserRole -from test_support.config import fast_config_validation - from splunk_ao.config import SplunkAOConfig from splunk_ao.utils.singleton import SplunkAOLoggerSingleton +from test_support.config import fast_config_validation # Note: The mock_request fixture is automatically provided by galileo_core[testing] extras diff --git a/splunk-ao-adk/tests/test_retriever.py b/splunk-ao-adk/tests/test_retriever.py index 13edab12..9b915030 100644 --- a/splunk-ao-adk/tests/test_retriever.py +++ b/splunk-ao-adk/tests/test_retriever.py @@ -9,7 +9,6 @@ from splunk_ao_adk.decorator import splunk_ao_retriever from splunk_ao_adk.observer import SplunkAOObserver - from .mocks import MockTool, MockToolContext diff --git a/splunk-ao-adk/tests/test_trace_builder.py b/splunk-ao-adk/tests/test_trace_builder.py index be97a9cf..c2d62299 100644 --- a/splunk-ao-adk/tests/test_trace_builder.py +++ b/splunk-ao-adk/tests/test_trace_builder.py @@ -3,8 +3,8 @@ from unittest.mock import MagicMock import pytest - from splunk_ao.schema.trace import TracesIngestRequest + from splunk_ao_adk.trace_builder import TraceBuilder diff --git a/src/splunk_ao/__init__.py b/src/splunk_ao/__init__.py index 15dd4e32..550421c8 100644 --- a/src/splunk_ao/__init__.py +++ b/src/splunk_ao/__init__.py @@ -19,7 +19,7 @@ from splunk_ao.collaborator import Collaborator, CollaboratorRole from splunk_ao.configuration import Configuration from splunk_ao.dataset import Dataset -from splunk_ao.decorator import SplunkAODecorator, log, splunk_ao_context, start_session +from splunk_ao.decorator import SplunkAODecorator, splunk_ao_context, log, start_session from splunk_ao.exceptions import ( AuthenticationError, BadRequestError, @@ -126,12 +126,12 @@ "create_api_key", "delete_api_key", "enable_console_logging", + "splunk_ao_context", "get_agent_control_target", "get_tracing_headers", "is_dependency_available", "list_api_keys", "log", "setup_agent_control_bridge", - "splunk_ao_context", "start_session", ] diff --git a/src/splunk_ao/constants/routes.py b/src/splunk_ao/constants/routes.py index 01ae0ab7..e2b6040d 100644 --- a/src/splunk_ao/constants/routes.py +++ b/src/splunk_ao/constants/routes.py @@ -1,7 +1,7 @@ -from enum import StrEnum +from enum import Enum -class Routes(StrEnum): +class Routes(str, Enum): healthcheck = "healthcheck" login = "login" api_key_login = "login/api_key" diff --git a/src/splunk_ao/experiment.py b/src/splunk_ao/experiment.py index 3f579818..07616946 100644 --- a/src/splunk_ao/experiment.py +++ b/src/splunk_ao/experiment.py @@ -429,7 +429,7 @@ def create(self) -> Experiment: if existing_experiment: _logger.warning(f"Experiment {existing_experiment.name} already exists, adding a timestamp") - now = datetime.datetime.now(datetime.UTC) + now = datetime.datetime.now(datetime.timezone.utc) self.name = f"{existing_experiment.name} {now:%Y-%m-%d} at {now:%H:%M:%S}.{now.microsecond // 1000:03d}" # Resolve prompt template before create (needed for trigger=True) diff --git a/src/splunk_ao/experiments.py b/src/splunk_ao/experiments.py index b2407552..4f5f1c0a 100644 --- a/src/splunk_ao/experiments.py +++ b/src/splunk_ao/experiments.py @@ -11,7 +11,7 @@ from galileo_core.constants.request_method import RequestMethod from splunk_ao.config import SplunkAOConfig from splunk_ao.datasets import Dataset, convert_dataset_row_to_record -from splunk_ao.decorator import log, splunk_ao_context, splunk_ao_dataset_context +from splunk_ao.decorator import splunk_ao_context, splunk_ao_dataset_context, log from splunk_ao.experiment_tags import upsert_experiment_tag from splunk_ao.projects import Project, Projects from splunk_ao.prompts import PromptTemplate @@ -298,9 +298,7 @@ def process_row(row: DatasetRecord, process_func: Callable) -> str: try: # Set dataset context for OTEL spans (ground truth for scorers) # This ensures OTEL-instrumented frameworks get dataset fields attached to their spans - with splunk_ao_dataset_context( - dataset_input=row.input, dataset_output=row.output, dataset_metadata=row.metadata - ): + with splunk_ao_dataset_context(dataset_input=row.input, dataset_output=row.output, dataset_metadata=row.metadata): output = process_func(row.deserialized_input) log = splunk_ao_context.get_logger_instance() log.conclude(output) @@ -426,7 +424,7 @@ def run_experiment( if existing_experiment: logging.warning(f"Experiment {existing_experiment.name} already exists, adding a timestamp") - now = datetime.datetime.now(datetime.UTC) + now = datetime.datetime.now(datetime.timezone.utc) experiment_name = f"{existing_experiment.name} {now:%Y-%m-%d} at {now:%H:%M:%S}.{now.microsecond // 1000:03d}" # Execute a runner function experiment (custom function flow — uses logstream pipeline) diff --git a/src/splunk_ao/handlers/agent_control/bridge.py b/src/splunk_ao/handlers/agent_control/bridge.py index e638983b..1da1190b 100644 --- a/src/splunk_ao/handlers/agent_control/bridge.py +++ b/src/splunk_ao/handlers/agent_control/bridge.py @@ -5,7 +5,7 @@ import threading import uuid from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import datetime, timezone from types import ModuleType from typing import Any @@ -53,9 +53,9 @@ def _load_agent_control_modules() -> _AgentControlModules: def _normalize_datetime(value: Any) -> datetime: if isinstance(value, datetime): if value.tzinfo is None: - return value.replace(tzinfo=UTC) + return value.replace(tzinfo=timezone.utc) return value - return datetime.now(tz=UTC) + return datetime.now(tz=timezone.utc) def _duration_ms_to_ns(value: Any) -> int | None: diff --git a/src/splunk_ao/handlers/base_handler.py b/src/splunk_ao/handlers/base_handler.py index ccdaf267..65a22f91 100644 --- a/src/splunk_ao/handlers/base_handler.py +++ b/src/splunk_ao/handlers/base_handler.py @@ -1,7 +1,7 @@ import logging import time from collections.abc import Callable -from datetime import UTC, datetime +from datetime import datetime, timezone from typing import Any from uuid import UUID @@ -257,7 +257,7 @@ def start_node(self, node_type: NODE_TYPE, parent_run_id: UUID | None, run_id: U node.span_params["start_time"] = time.perf_counter_ns() if "created_at" not in node.span_params: - node.span_params["created_at"] = datetime.now(tz=UTC) + node.span_params["created_at"] = datetime.now(tz=timezone.utc) found_node = self._nodes.get(node_id) if found_node: diff --git a/src/splunk_ao/handlers/openai_agents/handler.py b/src/splunk_ao/handlers/openai_agents/handler.py index 415203ea..e34a8032 100644 --- a/src/splunk_ao/handlers/openai_agents/handler.py +++ b/src/splunk_ao/handlers/openai_agents/handler.py @@ -1,6 +1,6 @@ import logging import uuid -from datetime import UTC, datetime +from datetime import datetime, timezone from typing import Any, cast from agents import Span, Trace, TracingProcessor @@ -66,7 +66,7 @@ def on_trace_start(self, trace: Trace) -> None: run_id=trace.trace_id, span_params={ "start_time": _get_timestamp(), - "start_time_iso": datetime.now(UTC).isoformat(), + "start_time_iso": datetime.now(timezone.utc).isoformat(), "name": trace.name, "metadata": convert_to_string_dict(trace.metadata), }, @@ -242,7 +242,7 @@ def on_span_start(self, span: Span[Any]) -> None: # Extract initial data based on type initial_params: dict[str, Any] = { "name": span_name, - "start_time_iso": span.started_at or datetime.now(UTC).isoformat(), + "start_time_iso": span.started_at or datetime.now(timezone.utc).isoformat(), } if splunk_ao_type in ["llm", "chat"]: llm_data = _extract_llm_data(span.span_data) @@ -316,7 +316,7 @@ def on_span_end(self, span: Span[Any]) -> None: # Update node with final data splunk_ao_type = node.node_type - end_params: dict[str, Any] = {"end_time_iso": span.ended_at or datetime.now(UTC).isoformat()} + end_params: dict[str, Any] = {"end_time_iso": span.ended_at or datetime.now(timezone.utc).isoformat()} end_params["duration_ns"] = convert_time_delta_to_ns( datetime.fromisoformat(span.ended_at) - datetime.fromisoformat(node.span_params["start_time_iso"]) ) diff --git a/src/splunk_ao/logger/control.py b/src/splunk_ao/logger/control.py index 59cef78f..34893aff 100644 --- a/src/splunk_ao/logger/control.py +++ b/src/splunk_ao/logger/control.py @@ -1,7 +1,7 @@ from __future__ import annotations -from datetime import UTC, datetime -from enum import StrEnum +from datetime import datetime, timezone +from enum import Enum from typing import Any, Literal from uuid import UUID @@ -17,15 +17,15 @@ except ImportError: HAS_NATIVE_CONTROL_SPAN = False - class ControlAppliesTo(StrEnum): + class ControlAppliesTo(str, Enum): llm_call = "llm_call" tool_call = "tool_call" - class ControlCheckStage(StrEnum): + class ControlCheckStage(str, Enum): pre = "pre" post = "post" - class ControlAction(StrEnum): + class ControlAction(str, Enum): deny = "deny" steer = "steer" observe = "observe" @@ -55,7 +55,7 @@ class ControlSpan(BaseModel): ) name: str = Field(default="control", description="Human-readable control name.") created_at: datetime = Field( - default_factory=lambda: datetime.now(tz=UTC), description="Timestamp of the control execution." + default_factory=lambda: datetime.now(tz=timezone.utc), description="Timestamp of the control execution." ) user_metadata: dict[str, str] = Field(default_factory=dict, description="Metadata associated with the span.") tags: list[str] = Field(default_factory=list, description="Tags associated with the span.") diff --git a/src/splunk_ao/schema/handlers.py b/src/splunk_ao/schema/handlers.py index a8903ecb..b601633c 100644 --- a/src/splunk_ao/schema/handlers.py +++ b/src/splunk_ao/schema/handlers.py @@ -1,4 +1,4 @@ -from enum import StrEnum +from enum import Enum from typing import Any, Literal from uuid import UUID @@ -8,7 +8,7 @@ INTEGRATION = Literal["langchain", "crewai", "google_adk"] -class NodeType(StrEnum): +class NodeType(str, Enum): AGENT = "agent" CHAIN = "chain" CHAT = "chat" diff --git a/src/splunk_ao/schema/message.py b/src/splunk_ao/schema/message.py index cf1b519e..bb183ee6 100644 --- a/src/splunk_ao/schema/message.py +++ b/src/splunk_ao/schema/message.py @@ -23,9 +23,6 @@ def __eq__(self, value: CoreMessage) -> bool: and self.tool_call_id == value.tool_call_id ) - def __hash__(self) -> int: - return hash((self.content, self.role, self.tool_call_id)) - # Without rebuilding the model, Message class we create here would not know and validate # constituent classes defined in core which build up message. diff --git a/src/splunk_ao/schema/metrics.py b/src/splunk_ao/schema/metrics.py index 47bf3b6b..27e25025 100644 --- a/src/splunk_ao/schema/metrics.py +++ b/src/splunk_ao/schema/metrics.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from enum import StrEnum +from enum import Enum from typing import Any, Generic, TypeVar from pydantic import BaseModel, Field, ValidationError, field_validator @@ -11,7 +11,7 @@ from galileo_core.schemas.shared.metric import MetricValueType -class SplunkAOMetrics(StrEnum): +class SplunkAOMetrics(str, Enum): """Built-in Galileo metric scorers. Values are human-readable UI labels used for scorer lookup via the API. diff --git a/src/splunk_ao/schema/trace.py b/src/splunk_ao/schema/trace.py index 234f365b..f3be6cc4 100644 --- a/src/splunk_ao/schema/trace.py +++ b/src/splunk_ao/schema/trace.py @@ -1,4 +1,4 @@ -from enum import StrEnum +from enum import Enum from typing import Any, Literal from pydantic import UUID4, BaseModel, Field @@ -10,7 +10,7 @@ SPAN_TYPE = Literal["llm", "retriever", "tool", "workflow", "agent"] -class LoggingMethod(StrEnum): +class LoggingMethod(str, Enum): playground = "playground" python_client = "python_client" typescript_client = "typescript_client" @@ -118,7 +118,7 @@ class SessionCreateResponse(BaseLogStreamOrExperimentModel): log_stream_id: UUID4 | None = Field(default=None, description="Log stream id associated with the session.") -class LogRecordsSearchFilterOperator(StrEnum): +class LogRecordsSearchFilterOperator(str, Enum): eq = "eq" ne = "ne" contains = "contains" @@ -131,7 +131,7 @@ class LogRecordsSearchFilterOperator(StrEnum): between = "between" -class LogRecordsSearchFilterType(StrEnum): +class LogRecordsSearchFilterType(str, Enum): id = "id" date = "date" number = "number" diff --git a/src/splunk_ao/search.py b/src/splunk_ao/search.py index b79cdfc4..0eb3c65f 100644 --- a/src/splunk_ao/search.py +++ b/src/splunk_ao/search.py @@ -1,5 +1,5 @@ import logging -from enum import StrEnum +from enum import Enum from splunk_ao.config import SplunkAOConfig from splunk_ao.resources.api.trace import ( @@ -18,7 +18,7 @@ logger = logging.getLogger(__name__) -class RecordType(StrEnum): +class RecordType(str, Enum): SPAN = "spans" TRACE = "traces" SESSION = "sessions" diff --git a/src/splunk_ao/utils/__init__.py b/src/splunk_ao/utils/__init__.py index cf16c651..4cc0e5fe 100644 --- a/src/splunk_ao/utils/__init__.py +++ b/src/splunk_ao/utils/__init__.py @@ -1,8 +1,8 @@ -from datetime import UTC, datetime +from datetime import datetime, timezone def _get_timestamp() -> datetime: - return datetime.now(UTC) + return datetime.now(timezone.utc) def _now_ns() -> int: diff --git a/src/splunk_ao/utils/serialization.py b/src/splunk_ao/utils/serialization.py index 27dead39..4cf69a2e 100644 --- a/src/splunk_ao/utils/serialization.py +++ b/src/splunk_ao/utils/serialization.py @@ -112,7 +112,7 @@ def serialize_datetime(v: dt.datetime) -> str: """ def _serialize_zoned_datetime(v: dt.datetime) -> str: - if v.tzinfo is not None and v.tzinfo.tzname(None) == dt.UTC.tzname(None): + if v.tzinfo is not None and v.tzinfo.tzname(None) == dt.timezone.utc.tzname(None): # UTC is a special case where we use "Z" at the end instead of "+00:00" return v.isoformat().replace("+00:00", "Z") # Delegate to the typical +/- offset format diff --git a/tests/test_agent_control_bridge.py b/tests/test_agent_control_bridge.py index 71f76e23..116c6fde 100644 --- a/tests/test_agent_control_bridge.py +++ b/tests/test_agent_control_bridge.py @@ -129,7 +129,7 @@ def _make_event(logger: SplunkAOLogger, **overrides: object) -> FakeControlExecu "action": "observe", "matched": True, "confidence": 0.91, - "timestamp": datetime.datetime.now(tz=datetime.UTC), + "timestamp": datetime.datetime.now(tz=datetime.timezone.utc), "execution_duration_ms": 12.5, "evaluator_name": "regex", "selector_path": "input", diff --git a/tests/test_base_handler.py b/tests/test_base_handler.py index 80348a1a..79eaadcf 100644 --- a/tests/test_base_handler.py +++ b/tests/test_base_handler.py @@ -38,9 +38,7 @@ def test_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: assert handler._nodes == {} # Custom initialization - handler = SplunkAOBaseHandler( - splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False - ) + handler = SplunkAOBaseHandler(splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False) assert handler._start_new_trace is False assert handler._flush_on_chain_end is False diff --git a/tests/test_config.py b/tests/test_config.py index b3efde3f..a4ac60a6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -71,7 +71,9 @@ def test_bridge_env_vars_propagates_splunk_ao_to_galileo(splunk_key, galileo_key with patch.dict(os.environ, {splunk_key: value}, clear=False): os.environ.pop(galileo_key, None) SplunkAOConfig._bridge_env_vars() - assert os.environ.get(galileo_key) == value, f"Expected {galileo_key}={value!r} after bridging {splunk_key}" + assert os.environ.get(galileo_key) == value, ( + f"Expected {galileo_key}={value!r} after bridging {splunk_key}" + ) @pytest.mark.parametrize("splunk_key,galileo_key", _ALL_BRIDGE_PAIRS) @@ -95,7 +97,9 @@ def test_bridge_env_vars_skips_absent_splunk_ao_keys() -> None: with patch.dict(os.environ, clean_env, clear=True): SplunkAOConfig._bridge_env_vars() for _, galileo_key in _ALL_BRIDGE_PAIRS: - assert galileo_key not in os.environ, f"{galileo_key} must not be set when its SPLUNK_AO_* source is absent" + assert galileo_key not in os.environ, ( + f"{galileo_key} must not be set when its SPLUNK_AO_* source is absent" + ) # --------------------------------------------------------------------------- diff --git a/tests/test_configuration.py b/tests/test_configuration.py index eb182f75..4eeb974a 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -244,7 +244,7 @@ def test_malformed_env_file_does_not_crash( capture_logs: tuple[logging.Logger, StringIO], ) -> None: """Test that malformed .env file logs error but doesn't crash the application.""" - _, _log_stream = capture_logs + _, log_stream = capture_logs # Create malformed env file (will cause parsing to fail) mock_env_file.write_text("INVALID_LINE_WITHOUT_EQUALS\n") @@ -313,7 +313,7 @@ def test_connect_fails_without_api_key( capture_logs: tuple[logging.Logger, StringIO], ) -> None: """Test connect() raises ConfigurationError when API key is missing.""" - _, _log_stream = capture_logs + _, log_stream = capture_logs # Set only console URL monkeypatch.setenv("SPLUNK_AO_CONSOLE_URL", "https://app.splunkao.ai") diff --git a/tests/test_crewai_handler.py b/tests/test_crewai_handler.py index e97a292d..c525c767 100644 --- a/tests/test_crewai_handler.py +++ b/tests/test_crewai_handler.py @@ -9,9 +9,13 @@ # Skip all tests in this module on Python 3.14+ (crewai doesn't support it yet) pytestmark = pytest.mark.skipif(sys.version_info >= (3, 14), reason="crewai does not support Python 3.14+") -from splunk_ao.handlers.crewai.handler import CrewAIEventListener -from splunk_ao.schema.handlers import NodeType -from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +from splunk_ao.handlers.crewai.handler import CrewAIEventListener # noqa: E402 +from splunk_ao.schema.handlers import NodeType # noqa: E402 +from tests.testutils.setup import ( # noqa: E402 + setup_mock_logstreams_client, + setup_mock_projects_client, + setup_mock_traces_client, +) class MockEvent: @@ -206,7 +210,7 @@ def test_extract_metadata(crewai_callback) -> None: assert metadata["key2"] == "value2" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_crew_kickoff_started(crewai_callback, generated_id) -> None: """Test crew kickoff started event handling.""" crew_id = generated_id() @@ -224,7 +228,7 @@ def test_crew_kickoff_started(crewai_callback, generated_id) -> None: assert call_args[1]["name"] == "Test Crew" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_crew_kickoff_started_empty_inputs(crewai_callback, generated_id) -> None: """Test crew kickoff started event handling.""" crew_id = generated_id() @@ -243,7 +247,7 @@ def test_crew_kickoff_started_empty_inputs(crewai_callback, generated_id) -> Non assert call_args[1]["input"] == "-" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_crew_kickoff_completed(crewai_callback, generated_id) -> None: """Test crew kickoff completed event handling.""" crew_id = generated_id() @@ -261,7 +265,7 @@ def test_crew_kickoff_completed(crewai_callback, generated_id) -> None: mock_end_node.assert_called_once_with(run_id=crew_id, output="Crew completed successfully") -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_crew_kickoff_failed(crewai_callback, generated_id) -> None: """Test crew kickoff failed event handling.""" crew_id = generated_id() @@ -278,7 +282,7 @@ def test_crew_kickoff_failed(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["error"] == "Something went wrong" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_agent_execution_started(crewai_callback, generated_id) -> None: """Test agent execution started event handling.""" agent_id = generated_id() @@ -301,7 +305,7 @@ def test_agent_execution_started(crewai_callback, generated_id) -> None: assert call_args[1]["input"] == "Research the topic" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_agent_execution_started_no_input(crewai_callback, generated_id) -> None: """Test agent execution started event handling.""" agent_id = generated_id() @@ -324,7 +328,7 @@ def test_agent_execution_started_no_input(crewai_callback, generated_id) -> None assert call_args[1]["input"] == "-" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_agent_execution_completed(crewai_callback, generated_id) -> None: """Test agent execution completed event handling.""" agent_id = generated_id() @@ -337,7 +341,7 @@ def test_agent_execution_completed(crewai_callback, generated_id) -> None: mock_end_node.assert_called_once_with(run_id=agent_id, output="Agent task completed") -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_agent_execution_error(crewai_callback, generated_id) -> None: """Test agent execution error event handling.""" agent_id = generated_id() @@ -354,7 +358,7 @@ def test_agent_execution_error(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["error"] == "Agent failed" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_task_started(crewai_callback, generated_id) -> None: """Test task started event handling.""" task_id = generated_id() @@ -377,7 +381,7 @@ def test_task_started(crewai_callback, generated_id) -> None: assert call_args[1]["input"] == "Previous research context" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_task_started_no_context(crewai_callback, generated_id) -> None: """Test task started event handling.""" task_id = generated_id() @@ -400,7 +404,7 @@ def test_task_started_no_context(crewai_callback, generated_id) -> None: assert call_args[1]["input"] == task.description -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_task_started_no_description(crewai_callback, generated_id) -> None: """Test task started event handling.""" task_id = generated_id() @@ -423,7 +427,7 @@ def test_task_started_no_description(crewai_callback, generated_id) -> None: assert call_args[1]["input"] == "-" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_task_completed(crewai_callback, generated_id) -> None: """Test task completed event handling.""" task_id = generated_id() @@ -437,7 +441,7 @@ def test_task_completed(crewai_callback, generated_id) -> None: mock_end_node.assert_called_once_with(run_id=task_id, output="Task completed successfully") -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_task_failed(crewai_callback, generated_id) -> None: """Test task failed event handling.""" task_id = generated_id() @@ -453,7 +457,7 @@ def test_task_failed(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["error"] == "Task execution failed" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_tool_usage_started(crewai_callback, generated_id) -> None: """Test tool usage started event handling.""" tool_id = generated_id() @@ -473,7 +477,7 @@ def test_tool_usage_started(crewai_callback, generated_id) -> None: assert call_args[1]["name"] == "search_tool" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_tool_usage_started_no_input(crewai_callback, generated_id) -> None: """Test tool usage started event handling.""" tool_id = generated_id() @@ -495,7 +499,7 @@ def test_tool_usage_started_no_input(crewai_callback, generated_id) -> None: assert call_args[1]["input"] == "-" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_tool_usage_finished(crewai_callback, generated_id) -> None: """Test tool usage finished event handling.""" tool_id = generated_id() @@ -508,7 +512,7 @@ def test_tool_usage_finished(crewai_callback, generated_id) -> None: mock_end_node.assert_called_once_with(run_id=tool_id, output="Tool execution completed") -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_tool_usage_error(crewai_callback, generated_id) -> None: """Test tool usage error event handling.""" tool_id = generated_id() @@ -525,7 +529,7 @@ def test_tool_usage_error(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["error"] == "Tool execution failed" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_llm_call_started(crewai_callback, generated_id) -> None: """Test LLM call started event handling.""" llm_id = generated_id() @@ -544,7 +548,7 @@ def test_llm_call_started(crewai_callback, generated_id) -> None: assert call_args[1]["temperature"] == 0.7 -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_llm_call_completed(crewai_callback, generated_id) -> None: """Test LLM call completed event handling.""" llm_id = generated_id() @@ -619,7 +623,7 @@ def test_llm_call_completed_extracts_token_usage_from_source(crewai_callback) -> assert call_args["total_tokens"] == 100 -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_llm_call_failed(crewai_callback, generated_id) -> None: """Test LLM call failed event handling.""" llm_id = generated_id() @@ -719,7 +723,7 @@ def test_lite_llm_usage_callback_no_node(crewai_callback) -> None: # Memory event tests (for CrewAI >= 0.177.0) -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_query_started(crewai_callback, generated_id) -> None: """Test memory query started event handling.""" query_id = generated_id() @@ -751,7 +755,7 @@ def test_memory_query_started(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["agent_role"] == "Research Agent" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_query_completed(crewai_callback, generated_id) -> None: """Test memory query completed event handling.""" query_id = generated_id() @@ -770,7 +774,7 @@ def test_memory_query_completed(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["results_count"] == 2 -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_query_failed(crewai_callback, generated_id) -> None: """Test memory query failed event handling.""" query_id = generated_id() @@ -788,7 +792,7 @@ def test_memory_query_failed(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["error"] == "Connection timeout" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_save_started(crewai_callback, generated_id) -> None: """Test memory save started event handling.""" save_id = generated_id() @@ -816,7 +820,7 @@ def test_memory_save_started(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["agent_role"] == "Research Agent" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_save_started_no_value(crewai_callback, generated_id) -> None: """Test memory save started event handling with no value.""" save_id = generated_id() @@ -831,7 +835,7 @@ def test_memory_save_started_no_value(crewai_callback, generated_id) -> None: assert call_args[1]["input"] == "Memory content" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_save_completed(crewai_callback, generated_id) -> None: """Test memory save completed event handling.""" save_id = generated_id() @@ -848,7 +852,7 @@ def test_memory_save_completed(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["save_time_ms"] == 75.2 -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_save_failed(crewai_callback, generated_id) -> None: """Test memory save failed event handling.""" save_id = generated_id() @@ -865,7 +869,7 @@ def test_memory_save_failed(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["error"] == "Storage full" -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_retrieval_started(crewai_callback, generated_id) -> None: """Test memory retrieval started event handling.""" retrieval_id = generated_id() @@ -888,7 +892,7 @@ def test_memory_retrieval_started(crewai_callback, generated_id) -> None: assert call_args[1]["metadata"]["task_id"] == task_id -@pytest.mark.parametrize("generated_id", [uuid.uuid4, lambda: str(uuid.uuid4())]) +@pytest.mark.parametrize("generated_id", [lambda: uuid.uuid4(), lambda: str(uuid.uuid4())]) def test_memory_retrieval_completed(crewai_callback, generated_id) -> None: """Test memory retrieval completed event handling.""" retrieval_id = generated_id() diff --git a/tests/test_decorator.py b/tests/test_decorator.py index 6954c98b..d9c9e148 100644 --- a/tests/test_decorator.py +++ b/tests/test_decorator.py @@ -8,7 +8,7 @@ from galileo_core.schemas.logging.span import AgentSpan, LlmSpan, RetrieverSpan, ToolSpan, WorkflowSpan from galileo_core.schemas.shared.document import Document from galileo_core.schemas.shared.multimodal import ContentModality -from splunk_ao import Message, MessageRole, log, splunk_ao_context, start_session +from splunk_ao import Message, MessageRole, splunk_ao_context, log, start_session from splunk_ao.decorator import _session_id_context from splunk_ao.schema.content_blocks import DataContentBlock, TextContentBlock from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client diff --git a/tests/test_decorator_distributed.py b/tests/test_decorator_distributed.py index 7dea065a..a14baaf4 100644 --- a/tests/test_decorator_distributed.py +++ b/tests/test_decorator_distributed.py @@ -8,7 +8,7 @@ from galileo_core.schemas.shared.document import Document from galileo_core.schemas.shared.multimodal import ContentModality -from splunk_ao import Message, MessageRole, log, splunk_ao_context +from splunk_ao import Message, MessageRole, splunk_ao_context, log from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER from splunk_ao.decorator import _parent_id_context, _trace_id_context from splunk_ao.schema.content_blocks import DataContentBlock, TextContentBlock diff --git a/tests/test_experiment.py b/tests/test_experiment.py index 5929292a..83ff8860 100644 --- a/tests/test_experiment.py +++ b/tests/test_experiment.py @@ -1,5 +1,5 @@ import logging -from datetime import UTC, datetime +from datetime import datetime, timezone from unittest.mock import MagicMock, patch from uuid import uuid4 @@ -26,8 +26,8 @@ def mock_experiment_response() -> MagicMock: mock_response = MagicMock(spec=ExperimentResponse) mock_response.id = str(uuid4()) mock_response.name = "Test Experiment" - mock_response.created_at = datetime.now(UTC) - mock_response.updated_at = datetime.now(UTC) + mock_response.created_at = datetime.now(timezone.utc) + mock_response.updated_at = datetime.now(timezone.utc) mock_response.additional_properties = {} mock_response.dataset = None mock_response.prompt = None @@ -809,15 +809,15 @@ def test_list_retrieves_all_experiments( mock_exp1 = MagicMock(spec=ExperimentResponse) mock_exp1.id = str(uuid4()) mock_exp1.name = "Experiment 1" - mock_exp1.created_at = datetime.now(UTC) - mock_exp1.updated_at = datetime.now(UTC) + mock_exp1.created_at = datetime.now(timezone.utc) + mock_exp1.updated_at = datetime.now(timezone.utc) mock_exp1.additional_properties = {} mock_exp2 = MagicMock(spec=ExperimentResponse) mock_exp2.id = str(uuid4()) mock_exp2.name = "Experiment 2" - mock_exp2.created_at = datetime.now(UTC) - mock_exp2.updated_at = datetime.now(UTC) + mock_exp2.created_at = datetime.now(timezone.utc) + mock_exp2.updated_at = datetime.now(timezone.utc) mock_exp2.additional_properties = {} mock_experiments_service = MagicMock() @@ -1754,7 +1754,7 @@ def test_repr_returns_detailed_string(self, reset_configuration: None) -> None: experiment.name = "Test Experiment" experiment.id = "test-id" experiment.project_id = "project-id" - experiment.created_at = datetime.now(UTC) + experiment.created_at = datetime.now(timezone.utc) result = repr(experiment) diff --git a/tests/test_experiment_tags.py b/tests/test_experiment_tags.py index d2eaa4a9..b503336a 100644 --- a/tests/test_experiment_tags.py +++ b/tests/test_experiment_tags.py @@ -26,8 +26,8 @@ def sample_run_tag(): key="environment", value="production", tag_type="generic", - created_at=datetime.datetime.now(datetime.UTC), - updated_at=datetime.datetime.now(datetime.UTC), + created_at=datetime.datetime.now(datetime.timezone.utc), + updated_at=datetime.datetime.now(datetime.timezone.utc), created_by="test_user", ) diff --git a/tests/test_experiments.py b/tests/test_experiments.py index e0df3679..05729f81 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -814,16 +814,19 @@ def test_run_experiment_no_prompt_no_dataset_raises( name="length", scorer_fn=lambda step: len(step.input), scorable_types=["workflow"], - aggregator_fn=sum, + aggregator_fn=lambda lengths: sum(lengths), ), LocalMetricConfig[str]( name="output", scorer_fn=lambda step: step.output, scorable_types=["workflow"], - aggregator_fn=",".join, + aggregator_fn=lambda outputs: ",".join(outputs), ), LocalMetricConfig[float]( - name="decimal", scorer_fn=lambda step: 4.53, scorable_types=["workflow"], aggregator_fn=mean + name="decimal", + scorer_fn=lambda step: 4.53, + scorable_types=["workflow"], + aggregator_fn=lambda values: mean(values), ), LocalMetricConfig[bool]( name="bool", @@ -842,10 +845,14 @@ def test_run_experiment_no_prompt_no_dataset_raises( complex_trace_function, [ LocalMetricConfig[int]( - name="length", scorer_fn=lambda step: len(step.input[0].content), aggregator_fn=sum + name="length", + scorer_fn=lambda step: len(step.input[0].content), + aggregator_fn=lambda lengths: sum(lengths), ), LocalMetricConfig[str]( - name="output", scorer_fn=lambda step: step.output.content, aggregator_fn=",".join + name="output", + scorer_fn=lambda step: step.output.content, + aggregator_fn=lambda outputs: ",".join(outputs), ), ], 2, @@ -1186,9 +1193,9 @@ def test_run_experiment_with_runner_and_dataset( # Return dataset_content on first call (starting_token=0), then None to signal end of pagination mock_get_dataset_instance = mock_get_dataset.return_value mock_get_dataset_instance.get_content = MagicMock( - side_effect=lambda starting_token=0, limit=1000: ( - dataset_content_with_question if starting_token == 0 else None - ) + side_effect=lambda starting_token=0, limit=1000: dataset_content_with_question + if starting_token == 0 + else None ) def runner(input) -> str: @@ -1416,7 +1423,7 @@ def test_create_scorer_configs(self, mock_scorer_settings_class, mock_scorers_cl from splunk_ao.utils.metrics import create_metric_configs scorers, local_scorers = create_metric_configs( - "project_id", "experiment_id", ["metric1", LocalMetricConfig(name="length", scorer_fn=len)] + "project_id", "experiment_id", ["metric1", LocalMetricConfig(name="length", scorer_fn=lambda x: len(x))] ) assert len(scorers) == 1 # Should return one valid scorer assert len(local_scorers) == 1 # Should return one local scorer @@ -1476,7 +1483,7 @@ def test_create_scorer_configs_with_metric_objects(self, mock_scorer_settings_cl mock_scorers_instance.get_scorer_version.assert_called_once_with(scorer_id="3", version=2) # Test mixed input types - local_metric = LocalMetricConfig(name="length", scorer_fn=len) + local_metric = LocalMetricConfig(name="length", scorer_fn=lambda x: len(x)) from splunk_ao.utils.metrics import create_metric_configs diff --git a/tests/test_langchain.py b/tests/test_langchain.py index 25b57ce1..d97343f2 100644 --- a/tests/test_langchain.py +++ b/tests/test_langchain.py @@ -993,7 +993,7 @@ def test_on_chain_start_converts_uuid7(self, callback): ) mock_start_node.assert_called_once() - args, _kwargs = mock_start_node.call_args + args, kwargs = mock_start_node.call_args # Verify UUIDs were converted to UUID4 converted_parent_id = args[1] # parent_run_id @@ -1012,7 +1012,7 @@ def test_on_llm_new_token_converts_uuid7(self, callback): callback.on_llm_new_token(token="test", run_id=uuid7_run_id) mock_get_node.assert_called_once() - args, _kwargs = mock_get_node.call_args + args, kwargs = mock_get_node.call_args converted_run_id = args[0] # run_id assert converted_run_id.version == 4 @@ -1037,9 +1037,7 @@ def logger_mocks(self): [ lambda hook: SplunkAOCallback(ingestion_hook=hook), lambda hook: SplunkAOCallback(splunk_ao_logger=SplunkAOLogger(), ingestion_hook=hook), - lambda hook: SplunkAOCallback( - splunk_ao_logger=splunk_ao_context.get_logger_instance(), ingestion_hook=hook - ), + lambda hook: SplunkAOCallback(splunk_ao_logger=splunk_ao_context.get_logger_instance(), ingestion_hook=hook), ], ) def test_on_chain_end_with_ingestion_hook(self, callback_builder): diff --git a/tests/test_langchain_async.py b/tests/test_langchain_async.py index b819c6f2..ce69f066 100644 --- a/tests/test_langchain_async.py +++ b/tests/test_langchain_async.py @@ -47,9 +47,7 @@ async def test_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: assert callback._handler._nodes == {} # Custom initialization - callback = SplunkAOAsyncCallback( - splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False - ) + callback = SplunkAOAsyncCallback(splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False) assert callback._handler._start_new_trace is False assert callback._handler._flush_on_chain_end is False @@ -664,9 +662,7 @@ async def test_callback_with_active_trace(self, splunk_ao_logger: SplunkAOLogger splunk_ao_logger.start_trace(input="test input") # Pass the active logger to the callback - callback = SplunkAOAsyncCallback( - splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False - ) + callback = SplunkAOAsyncCallback(splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False) # Start a chain (creates a workflow span) await callback._handler.async_start_node("chain", None, run_id, name="Test Chain", input='{"query": "test"}') @@ -997,7 +993,7 @@ async def test_async_on_chain_start_converts_uuid7(self, callback): ) mock_start_node.assert_called_once() - args, _kwargs = mock_start_node.call_args + args, kwargs = mock_start_node.call_args # Verify UUIDs were converted to UUID4 converted_parent_id = args[1] # parent_run_id @@ -1017,7 +1013,7 @@ async def test_async_on_tool_error_converts_uuid7(self, callback): await callback.on_tool_error(error=Exception("test"), run_id=uuid7_run_id) mock_end_node.assert_called_once() - args, _kwargs = mock_end_node.call_args + args, kwargs = mock_end_node.call_args converted_run_id = args[0] # run_id assert converted_run_id.version == 4 diff --git a/tests/test_langchain_middleware.py b/tests/test_langchain_middleware.py index b46ff721..fa68996f 100644 --- a/tests/test_langchain_middleware.py +++ b/tests/test_langchain_middleware.py @@ -136,9 +136,7 @@ def test_default_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: def test_custom_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: """Test middleware initialization with custom parameters.""" - middleware = SplunkAOMiddleware( - splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False - ) + middleware = SplunkAOMiddleware(splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False) assert middleware._handler._start_new_trace is False assert middleware._handler._flush_on_chain_end is False @@ -510,9 +508,7 @@ def logger_mocks(self): [ lambda hook: SplunkAOMiddleware(ingestion_hook=hook), lambda hook: SplunkAOMiddleware(splunk_ao_logger=SplunkAOLogger(), ingestion_hook=hook), - lambda hook: SplunkAOMiddleware( - splunk_ao_logger=splunk_ao_context.get_logger_instance(), ingestion_hook=hook - ), + lambda hook: SplunkAOMiddleware(splunk_ao_logger=splunk_ao_context.get_logger_instance(), ingestion_hook=hook), ], ) def test_ingestion_hook_called(self, middleware_builder) -> None: diff --git a/tests/test_log_streams_metrics.py b/tests/test_log_streams_metrics.py index 1d972a87..1beac71d 100644 --- a/tests/test_log_streams_metrics.py +++ b/tests/test_log_streams_metrics.py @@ -89,7 +89,7 @@ def test_create_metric_configs_with_builtin_metrics( mock_scorer_settings_class.return_value.create.return_value = None # Test with built-in metrics - _scorers, local_metrics = create_metric_configs( + scorers, local_metrics = create_metric_configs( "project-123", "logstream-456", [SplunkAOMetrics.correctness, "completeness"] ) @@ -142,7 +142,7 @@ def custom_scorer(trace_or_span) -> float: local_metric = LocalMetricConfig(name="local_metric", scorer_fn=custom_scorer) # Test with mixed metrics (only valid ones to avoid decorator error handling) - _scorers, local_metrics = create_metric_configs( + scorers, local_metrics = create_metric_configs( "project-123", "logstream-456", [SplunkAOMetrics.correctness, local_metric] ) diff --git a/tests/test_logger_batch.py b/tests/test_logger_batch.py index 768e7443..87f624e5 100644 --- a/tests/test_logger_batch.py +++ b/tests/test_logger_batch.py @@ -1179,7 +1179,7 @@ def test_get_last_output_llm_message_raw() -> None: trace.spans = [llm_span] # When: getting the last output - output, _redacted_output = SplunkAOLogger._get_last_output(trace) + output, redacted_output = SplunkAOLogger._get_last_output(trace) # Then: the raw Message is returned (caller coerces for Trace destinations) assert isinstance(output, Message) diff --git a/tests/test_logger_distributed.py b/tests/test_logger_distributed.py index ff651dc1..df479e8f 100644 --- a/tests/test_logger_distributed.py +++ b/tests/test_logger_distributed.py @@ -1,6 +1,8 @@ import asyncio import datetime +import json import logging +import uuid from unittest.mock import Mock, patch from uuid import UUID diff --git a/tests/test_logger_timestamps.py b/tests/test_logger_timestamps.py index dd61e242..56bae4a1 100644 --- a/tests/test_logger_timestamps.py +++ b/tests/test_logger_timestamps.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta +from datetime import datetime, timedelta, timezone from unittest.mock import Mock, patch from splunk_ao.logger import SplunkAOLogger @@ -33,7 +33,7 @@ def test_user_provided_timestamps_are_respected(mock_projects_client: Mock, mock logger.start_trace(input="test") # Create timestamps in reverse order to test that the logger doesn't alter them - now = datetime.now(UTC) + now = datetime.now(timezone.utc) timestamps = [now - timedelta(seconds=i) for i in range(5)] for ts in timestamps: @@ -59,11 +59,11 @@ def test_mixed_default_and_user_timestamps(mock_projects_client: Mock, mock_logs logger.add_llm_span(input="test", output="test", model="test") # 2. Add a user-provided span with a timestamp in the past - past_timestamp = datetime.now(UTC) - timedelta(seconds=10) + past_timestamp = datetime.now(timezone.utc) - timedelta(seconds=10) logger.add_llm_span(input="test", output="test", model="test", created_at=past_timestamp) # 3. Add a user-provided span with a timestamp in the future - future_timestamp = datetime.now(UTC) + timedelta(seconds=10) + future_timestamp = datetime.now(timezone.utc) + timedelta(seconds=10) logger.add_llm_span(input="test", output="test", model="test", created_at=future_timestamp) # 4. Add a final default span diff --git a/tests/test_openai.py b/tests/test_openai.py index 6f17bf23..4599a8f9 100644 --- a/tests/test_openai.py +++ b/tests/test_openai.py @@ -8,7 +8,7 @@ from openai.types.responses import ResponseCompletedEvent from galileo_core.schemas.logging.span import LlmSpan, WorkflowSpan -from splunk_ao import Message, MessageRole, log, splunk_ao_context +from splunk_ao import Message, MessageRole, splunk_ao_context, log from splunk_ao.openai import OpenAIGalileo, openai from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client from tests.testutils.streaming import EventStream, ResponsesEventStream diff --git a/tests/utils/test_datasets.py b/tests/utils/test_datasets.py index 69de2dd2..67accb8f 100644 --- a/tests/utils/test_datasets.py +++ b/tests/utils/test_datasets.py @@ -25,7 +25,7 @@ def test_get_dataset_and_records_with_id(mock_get_records, mock_get_dataset, dat mock_get_dataset.return_value = mock_dataset # Execute - dataset, _records = get_dataset_and_records(id="test-id") + dataset, records = get_dataset_and_records(id="test-id") # Assert mock_get_dataset.assert_called_once_with(id="test-id") @@ -43,7 +43,7 @@ def test_get_dataset_and_records_with_name(mock_get_records, mock_get_dataset, d mock_get_dataset.return_value = mock_dataset # Execute - dataset, _records = get_dataset_and_records(name="test-dataset") + dataset, records = get_dataset_and_records(name="test-dataset") # Assert mock_get_dataset.assert_called_once_with(name="test-dataset") diff --git a/tests/utils/test_serialization.py b/tests/utils/test_serialization.py index 561326e0..300b2ef7 100644 --- a/tests/utils/test_serialization.py +++ b/tests/utils/test_serialization.py @@ -27,7 +27,7 @@ class TestSerializeDateTime: def test_serialize_datetime_with_utc_timezone(self) -> None: # Test with UTC timezone - dt_utc = dt.datetime(2023, 1, 1, 12, 0, 0, tzinfo=dt.UTC) + dt_utc = dt.datetime(2023, 1, 1, 12, 0, 0, tzinfo=dt.timezone.utc) result = serialize_datetime(dt_utc) assert result.endswith("Z") assert result == "2023-01-01T12:00:00Z" @@ -44,7 +44,7 @@ def test_serialize_datetime_with_non_utc_timezone(self) -> None: class TestEventSerializer: def test_default_datetime(self) -> None: # Test datetime serialization - dt_obj = dt.datetime(2023, 1, 1, 12, 0, 0, tzinfo=dt.UTC) + dt_obj = dt.datetime(2023, 1, 1, 12, 0, 0, tzinfo=dt.timezone.utc) result = json.dumps(dt_obj, cls=EventSerializer) decoded_result = json.loads(result) assert decoded_result == "2023-01-01T12:00:00Z" From 19fc0c8944262f90f8755e72d008d99aed18d042 Mon Sep 17 00:00:00 2001 From: etserend Date: Mon, 13 Jul 2026 18:23:59 -0500 Subject: [PATCH 2/2] fix(deps): regenerate poetry.lock after revert of ruff/openapi-python-client bump Co-Authored-By: Claude Opus 4.7 --- poetry.lock | 37 ++++++++++++++++++++++++++----------- 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/poetry.lock b/poetry.lock index d9e511f8..a6302d83 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -454,7 +454,7 @@ files = [ {file = "cffi-1.17.1-cp39-cp39-win_amd64.whl", hash = "sha256:d016c76bdd850f3c626af19b0542c9677ba156e4ee4fccfdd7848803533ef662"}, {file = "cffi-1.17.1.tar.gz", hash = "sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or platform_python_implementation == \"PyPy\" or (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") and python_version <= \"3.13\"", test = "platform_python_implementation == \"PyPy\""} +markers = {main = "platform_python_implementation != \"PyPy\" and (extra == \"openai\" or extra == \"all\") or platform_python_implementation != \"PyPy\" and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") or (extra == \"langchain\" or extra == \"all\") and platform_python_implementation == \"PyPy\" or (extra == \"langchain\" or extra == \"all\") and (extra == \"openai\" or extra == \"all\") or (extra == \"langchain\" or extra == \"all\") and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation == \"PyPy\""} [package.dependencies] pycparser = "*" @@ -559,6 +559,7 @@ files = [ {file = "charset_normalizer-3.4.3-py3-none-any.whl", hash = "sha256:ce571ab16d890d23b5c278547ba694193a45011ff86a9162a71307ed9f86759a"}, {file = "charset_normalizer-3.4.3.tar.gz", hash = "sha256:6fce4b8500244f6fcb71465d4a4930d132ba9ab8e71a7859e6a5d59851068d14"}, ] +markers = {main = "python_version < \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or extra == \"crewai\")"} [[package]] name = "chromadb" @@ -900,6 +901,7 @@ files = [ {file = "distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2"}, {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, ] +markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"openai\") and python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\""} [[package]] name = "docstring-parser" @@ -908,11 +910,11 @@ description = "Parse Python docstrings in reST, Google and Numpydoc format" optional = false python-versions = ">=3.8" groups = ["main", "test"] -markers = {main = "(extra == \"crewai\" or extra == \"all\") and python_version <= \"3.13\""} files = [ {file = "docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708"}, {file = "docstring_parser-0.17.0.tar.gz", hash = "sha256:583de4a309722b3315439bb31d64ba3eebada841f2e2cee23b99df001434c912"}, ] +markers = {main = "(extra == \"crewai\" or extra == \"all\") and python_version <= \"3.13\""} [package.extras] dev = ["pre-commit (>=2.16.0) ; python_version >= \"3.9\"", "pydoctor (>=25.4.0)", "pytest"] @@ -1876,6 +1878,7 @@ files = [ {file = "jiter-0.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:1b28302349dc65703a9e4ead16f163b1c339efffbe1049c30a44b001a2a4fff9"}, {file = "jiter-0.10.0.tar.gz", hash = "sha256:07a7142c38aacc85194391108dc91b5b57093c978a9932bd86a36862759d9500"}, ] +markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"openai\") and python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\""} [[package]] name = "json-repair" @@ -1917,6 +1920,7 @@ files = [ {file = "jsonpatch-1.33-py2.py3-none-any.whl", hash = "sha256:0ae28c0cd062bbd8b8ecc26d7d164fbbea9652a1a3693f3b956c1eae5145dade"}, {file = "jsonpatch-1.33.tar.gz", hash = "sha256:9fcd4009c41e6d12348b4a0ff2563ba56a2923a7dfee731d004e212e1ee5030c"}, ] +markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] jsonpointer = ">=1.9" @@ -1932,6 +1936,7 @@ files = [ {file = "jsonpointer-3.0.0-py2.py3-none-any.whl", hash = "sha256:13e088adc14fca8b6aa8177c044e12701e6ad4b28ff10e65f2267a90109c9942"}, {file = "jsonpointer-3.0.0.tar.gz", hash = "sha256:2b2d729f2091522d61c3b31f82e11870f60b68f43fbc705cb76bf4b832af59ef"}, ] +markers = {main = "extra == \"langchain\" or extra == \"all\""} [[package]] name = "jsonref" @@ -1961,7 +1966,7 @@ files = [ [package.dependencies] attrs = ">=22.2.0" -jsonschema-specifications = ">=2023.03.6" +jsonschema-specifications = ">=2023.3.6" referencing = ">=0.28.4" rpds-py = ">=0.7.1" @@ -1999,7 +2004,7 @@ files = [ ] [package.dependencies] -certifi = ">=14.05.14" +certifi = ">=14.5.14" durationpy = ">=0.7" google-auth = ">=1.0.1" oauthlib = ">=3.2.2" @@ -2061,6 +2066,7 @@ files = [ {file = "langchain_core-1.2.7-py3-none-any.whl", hash = "sha256:452f4fef7a3d883357b22600788d37e3d8854ef29da345b7ac7099f33c31828b"}, {file = "langchain_core-1.2.7.tar.gz", hash = "sha256:e1460639f96c352b4a41c375f25aeb8d16ffc1769499fb1c20503aad59305ced"}, ] +markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] jsonpatch = ">=1.33.0,<2.0.0" @@ -2155,6 +2161,7 @@ files = [ {file = "langsmith-0.4.14-py3-none-any.whl", hash = "sha256:b6d070ac425196947d2a98126fb0e35f3b8c001a2e6e5b7049dd1c56f0767d0b"}, {file = "langsmith-0.4.14.tar.gz", hash = "sha256:4d29c7a9c85b20ba813ab9c855407bccdf5eb4f397f512ffa89959b2a2cb83ed"}, ] +markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] httpx = ">=0.23.0,<1" @@ -2852,6 +2859,7 @@ files = [ {file = "openai-2.32.0-py3-none-any.whl", hash = "sha256:4dcc9badeb4bf54ad0d187453742f290226d30150890b7890711bda4f32f192f"}, {file = "openai-2.32.0.tar.gz", hash = "sha256:c54b27a9e4cb8d51f0dd94972ffd1a04437efeb259a9e60d8922b8bd26fe55e0"}, ] +markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"openai\") and python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\""} [package.dependencies] anyio = ">=3.5.0,<5" @@ -3195,7 +3203,7 @@ files = [ {file = "orjson-3.11.2-cp39-cp39-win_amd64.whl", hash = "sha256:c9ec0cc0d4308cad1e38a1ee23b64567e2ff364c2a3fe3d6cbc69cf911c45712"}, {file = "orjson-3.11.2.tar.gz", hash = "sha256:91bdcf5e69a8fd8e8bdb3de32b31ff01d2bd60c1e8d5fe7d5afabdcf19920309"}, ] -markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"langchain\" or extra == \"all\" or (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\") and python_version <= \"3.13\"", test = "platform_python_implementation != \"PyPy\""} +markers = {main = "extra == \"langchain\" or extra == \"all\" or (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") and python_version <= \"3.13\"", test = "platform_python_implementation != \"PyPy\""} [[package]] name = "ormsgpack" @@ -3280,6 +3288,7 @@ files = [ {file = "packaging-24.2-py3-none-any.whl", hash = "sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759"}, {file = "packaging-24.2.tar.gz", hash = "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f"}, ] +markers = {main = "python_version < \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"crewai\")"} [[package]] name = "pathspec" @@ -3945,7 +3954,7 @@ files = [ {file = "pycparser-2.22-py3-none-any.whl", hash = "sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc"}, {file = "pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or platform_python_implementation == \"PyPy\" or (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") and python_version <= \"3.13\"", test = "platform_python_implementation == \"PyPy\""} +markers = {main = "platform_python_implementation != \"PyPy\" and (extra == \"openai\" or extra == \"all\") or platform_python_implementation != \"PyPy\" and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\") or (extra == \"langchain\" or extra == \"all\") and platform_python_implementation == \"PyPy\" or (extra == \"langchain\" or extra == \"all\") and (extra == \"openai\" or extra == \"all\") or (extra == \"langchain\" or extra == \"all\") and python_version <= \"3.13\" and (extra == \"openai\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation == \"PyPy\""} [[package]] name = "pydantic" @@ -4530,6 +4539,7 @@ files = [ {file = "PyYAML-6.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8"}, {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"}, ] +markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") and python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\""} [[package]] name = "referencing" @@ -4658,6 +4668,7 @@ files = [ {file = "requests-2.32.4-py3-none-any.whl", hash = "sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c"}, {file = "requests-2.32.4.tar.gz", hash = "sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422"}, ] +markers = {main = "python_version < \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\" or extra == \"crewai\")"} [package.dependencies] certifi = ">=2017.4.17" @@ -4718,6 +4729,7 @@ files = [ {file = "requests-toolbelt-1.0.0.tar.gz", hash = "sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6"}, {file = "requests_toolbelt-1.0.0-py2.py3-none-any.whl", hash = "sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06"}, ] +markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] requests = ">=2.0.1,<3.0.0" @@ -5151,6 +5163,7 @@ files = [ {file = "tenacity-9.1.2-py3-none-any.whl", hash = "sha256:f77bf36710d8b73a50b2dd155c97b870017ad21afe6ab300326b0371b3b05138"}, {file = "tenacity-9.1.2.tar.gz", hash = "sha256:1169d376c297e7de388d18b4481760d478b0e99a777cad3a9c86e556f4b697cb"}, ] +markers = {main = "(extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") and python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\""} [package.extras] doc = ["reno", "sphinx"] @@ -5547,11 +5560,11 @@ description = "HTTP library with thread-safe connection pooling, file post, and optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7" groups = ["main", "test"] -markers = "platform_python_implementation == \"PyPy\"" files = [ {file = "urllib3-1.26.20-py2.py3-none-any.whl", hash = "sha256:0ed14ccfbf1c30a9072c7ca157e4319b70d65f623e91e7b32fadb2853431016e"}, {file = "urllib3-1.26.20.tar.gz", hash = "sha256:40c2dc0c681e47eb8f90e7e27bf6ff7df2e677421fd46756da1161c39ca70d32"}, ] +markers = {main = "platform_python_implementation == \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\") or platform_python_implementation == \"PyPy\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") and python_version <= \"3.13\"", test = "platform_python_implementation == \"PyPy\""} [package.extras] brotli = ["brotli (==1.0.9) ; os_name != \"nt\" and python_version < \"3\" and platform_python_implementation == \"CPython\"", "brotli (>=1.0.9) ; python_version >= \"3\" and platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; (os_name != \"nt\" or python_version >= \"3\") and platform_python_implementation != \"CPython\"", "brotlipy (>=0.6.0) ; os_name == \"nt\" and python_version < \"3\""] @@ -5565,11 +5578,11 @@ description = "HTTP library with thread-safe connection pooling, file post, and optional = false python-versions = ">=3.9" groups = ["main", "test"] -markers = "platform_python_implementation != \"PyPy\"" files = [ {file = "urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc"}, {file = "urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760"}, ] +markers = {main = "platform_python_implementation != \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"otel\") or platform_python_implementation != \"PyPy\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"otel\") and python_version <= \"3.13\"", test = "platform_python_implementation != \"PyPy\""} [package.extras] brotli = ["brotli (>=1.0.9) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\""] @@ -5608,6 +5621,7 @@ files = [ {file = "uuid_utils-0.13.0-pp311-pypy311_pp73-manylinux_2_24_x86_64.whl", hash = "sha256:b7ccaa20e24c5f60f41a69ef571ed820737f9b0ade4cbeef56aaa8f80f5aa475"}, {file = "uuid_utils-0.13.0.tar.gz", hash = "sha256:4c17df6427a9e23a4cd7fb9ee1efb53b8abb078660b9bdb2524ca8595022dfe1"}, ] +markers = {main = "extra == \"langchain\" or extra == \"all\""} [[package]] name = "uv" @@ -5734,8 +5748,8 @@ files = [ [package.dependencies] PyYAML = "*" urllib3 = [ - {version = "<2", markers = "platform_python_implementation == \"PyPy\""}, {version = "*", markers = "platform_python_implementation != \"PyPy\" and python_version >= \"3.10\""}, + {version = "<2", markers = "platform_python_implementation == \"PyPy\""}, ] wrapt = "*" yarl = "*" @@ -6504,6 +6518,7 @@ files = [ {file = "zstandard-0.23.0-cp39-cp39-win_amd64.whl", hash = "sha256:f8346bfa098532bc1fb6c7ef06783e969d87a99dd1d2a5a18a892c1d7a643c58"}, {file = "zstandard-0.23.0.tar.gz", hash = "sha256:b2d8c62d08e7255f68f7a740bae85b3c9b8e5466baa9cbf7f57f1cde0ac6bc09"}, ] +markers = {main = "extra == \"langchain\" or extra == \"all\""} [package.dependencies] cffi = {version = ">=1.11", markers = "platform_python_implementation == \"PyPy\""} @@ -6522,4 +6537,4 @@ otel = ["grpcio", "opentelemetry-api", "opentelemetry-exporter-otlp", "opentelem [metadata] lock-version = "2.1" python-versions = "^3.11,<3.15" -content-hash = "f75a182570e1129b6b65927aac9d700eae02d507391e7d2d6850d8910e4d55bf" +content-hash = "bd8c012c6d1fa222ef146ed6956d1ff0c658ad50c544514a26e3234188825864"