agentbasis 0.1.0__py3-none-any.whl

agentbasis/llms/openai/chat.py

@@ -0,0 +1,235 @@
+from typing import Any, Generator, AsyncGenerator
+import functools
+import time
+from opentelemetry import trace
+from opentelemetry.trace import Status, StatusCode, Span
+
+from agentbasis.context import inject_context_to_span
+
+
+def _get_tracer():
+    """
+    Get the tracer lazily at runtime.
+    This ensures the tracer is retrieved after agentbasis.init() has configured the provider.
+    """
+    return trace.get_tracer("agentbasis.llms.openai")
+
+
+def _set_request_attributes(span: Span, model: str, messages: list, is_streaming: bool = False):
+    """
+    Set common request attributes on a span.
+    """
+    # Inject user/session context
+    inject_context_to_span(span)
+    
+    span.set_attribute("llm.system", "openai")
+    span.set_attribute("llm.request.model", model)
+    span.set_attribute("llm.request.messages", str(messages))
+    if is_streaming:
+        span.set_attribute("llm.request.streaming", True)
+
+
+def _set_response_attributes(span: Span, response):
+    """
+    Set common response attributes on a span (for non-streaming responses).
+    """
+    if response.choices:
+        content = response.choices[0].message.content
+        span.set_attribute("llm.response.content", str(content))
+
+    if response.usage:
+        span.set_attribute("llm.usage.prompt_tokens", response.usage.prompt_tokens)
+        span.set_attribute("llm.usage.completion_tokens", response.usage.completion_tokens)
+        span.set_attribute("llm.usage.total_tokens", response.usage.total_tokens)
+
+
+def _wrap_sync_stream(stream, span: Span, start_time: float) -> Generator:
+    """
+    Wrap a synchronous streaming response to track chunks and finalize span.
+    """
+    content_parts = []
+    chunk_count = 0
+    first_token_time = None
+    
+    try:
+        for chunk in stream:
+            chunk_count += 1
+            
+            # Track time to first token
+            if first_token_time is None:
+                first_token_time = time.time()
+                span.set_attribute("llm.response.first_token_ms", 
+                                   int((first_token_time - start_time) * 1000))
+            
+            # Extract content from chunk
+            if chunk.choices and len(chunk.choices) > 0:
+                delta = chunk.choices[0].delta
+                if hasattr(delta, 'content') and delta.content:
+                    content_parts.append(delta.content)
+            
+            yield chunk
+        
+        # Stream complete - finalize span
+        full_content = "".join(content_parts)
+        span.set_attribute("llm.response.content", full_content)
+        span.set_attribute("llm.response.chunk_count", chunk_count)
+        span.set_status(Status(StatusCode.OK))
+        
+    except Exception as e:
+        span.record_exception(e)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        raise
+    finally:
+        span.end()
+
+
+async def _wrap_async_stream(stream, span: Span, start_time: float) -> AsyncGenerator:
+    """
+    Wrap an asynchronous streaming response to track chunks and finalize span.
+    """
+    content_parts = []
+    chunk_count = 0
+    first_token_time = None
+    
+    try:
+        async for chunk in stream:
+            chunk_count += 1
+            
+            # Track time to first token
+            if first_token_time is None:
+                first_token_time = time.time()
+                span.set_attribute("llm.response.first_token_ms", 
+                                   int((first_token_time - start_time) * 1000))
+            
+            # Extract content from chunk
+            if chunk.choices and len(chunk.choices) > 0:
+                delta = chunk.choices[0].delta
+                if hasattr(delta, 'content') and delta.content:
+                    content_parts.append(delta.content)
+            
+            yield chunk
+        
+        # Stream complete - finalize span
+        full_content = "".join(content_parts)
+        span.set_attribute("llm.response.content", full_content)
+        span.set_attribute("llm.response.chunk_count", chunk_count)
+        span.set_status(Status(StatusCode.OK))
+        
+    except Exception as e:
+        span.record_exception(e)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        raise
+    finally:
+        span.end()
+
+
+def instrument_chat(openai_module: Any):
+    """
+    Instruments the synchronous OpenAI Chat Completions API with OpenTelemetry.
+    Handles both regular and streaming responses.
+    """
+    try:
+        from openai.resources.chat.completions import Completions
+    except ImportError:
+        return
+
+    original_create = Completions.create
+
+    @functools.wraps(original_create)
+    def wrapped_create(self, *args, **kwargs):
+        tracer = _get_tracer()
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+        is_streaming = kwargs.get("stream", False)
+
+        span_name = f"openai.chat.completions.create {model}"
+        
+        if is_streaming:
+            # For streaming, we need to manually manage the span lifecycle
+            span = tracer.start_span(span_name)
+            start_time = time.time()
+            _set_request_attributes(span, model, messages, is_streaming=True)
+            
+            try:
+                stream = original_create(self, *args, **kwargs)
+                # Return wrapped generator that will finalize span when exhausted
+                return _wrap_sync_stream(stream, span, start_time)
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.end()
+                raise
+        else:
+            # Non-streaming: use context manager as before
+            with tracer.start_as_current_span(span_name) as span:
+                _set_request_attributes(span, model, messages)
+
+                try:
+                    response = original_create(self, *args, **kwargs)
+                    _set_response_attributes(span, response)
+                    span.set_status(Status(StatusCode.OK))
+                    return response
+
+                except Exception as e:
+                    span.record_exception(e)
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    raise
+
+    Completions.create = wrapped_create
+
+
+def instrument_async_chat(openai_module: Any):
+    """
+    Instruments the asynchronous OpenAI Chat Completions API with OpenTelemetry.
+    Handles both regular and streaming responses.
+    """
+    try:
+        from openai.resources.chat.completions import AsyncCompletions
+    except ImportError:
+        return
+
+    original_async_create = AsyncCompletions.create
+
+    @functools.wraps(original_async_create)
+    async def wrapped_async_create(self, *args, **kwargs):
+        tracer = _get_tracer()
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+        is_streaming = kwargs.get("stream", False)
+
+        span_name = f"openai.chat.completions.create {model}"
+        
+        if is_streaming:
+            # For streaming, we need to manually manage the span lifecycle
+            span = tracer.start_span(span_name)
+            start_time = time.time()
+            span.set_attribute("llm.request.async", True)
+            _set_request_attributes(span, model, messages, is_streaming=True)
+            
+            try:
+                stream = await original_async_create(self, *args, **kwargs)
+                # Return wrapped async generator that will finalize span when exhausted
+                return _wrap_async_stream(stream, span, start_time)
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.end()
+                raise
+        else:
+            # Non-streaming: use context manager as before
+            with tracer.start_as_current_span(span_name) as span:
+                span.set_attribute("llm.request.async", True)
+                _set_request_attributes(span, model, messages)
+
+                try:
+                    response = await original_async_create(self, *args, **kwargs)
+                    _set_response_attributes(span, response)
+                    span.set_status(Status(StatusCode.OK))
+                    return response
+
+                except Exception as e:
+                    span.record_exception(e)
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    raise
+
+    AsyncCompletions.create = wrapped_async_create

agentbasis-0.1.0.dist-info/METADATA

@@ -0,0 +1,220 @@
+Metadata-Version: 2.4
+Name: agentbasis
+Version: 0.1.0
+Summary: Management & Observability SDK for AI Agents
+Author-email: AgentBasis <support@agentbasis.co>
+Maintainer-email: AgentBasis <support@agentbasis.co>
+Project-URL: Homepage, https://agentbasis.co
+Project-URL: Documentation, https://docs.agentbasis.co
+Project-URL: Repository, https://github.com/AgentBasis/agentbasis-python-sdk
+Project-URL: Bug Tracker, https://github.com/AgentBasis/agentbasis-python-sdk/issues
+Keywords: ai,agents,observability,tracing,opentelemetry,llm,openai,anthropic,langchain,monitoring,telemetry
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: System :: Monitoring
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: requests>=2.25.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: opentelemetry-api>=1.20.0
+Requires-Dist: opentelemetry-sdk>=1.20.0
+Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.20.0
+
+# AgentBasis Python SDK
+
+**Management & Observability SDK for AI Agents in production**
+
+The **AgentBasis Python SDK** provides a simple, lightweight way to track the performance, traces, sessions, and behavior of AI agents. It sends data using the **OpenTelemetry (OTel)** standard, making it compatible with AgentBasis and other observability backends.
+
+This is the **foundation SDK** that enables deep observability for coded agents built with:
+- Pure Python
+- LLM Providers: 
+    - OpenAI
+    - Anthropic
+    - Gemini
+- Frameworks
+    - LangChain
+    - Pydantic AI
+
+## Installation
+
+```bash
+pip install agentbasis
+```
+
+## Quick Start
+
+### 1. Initialize the SDK
+Start by initializing the SDK with your API key and Agent ID. This usually goes at the top of your main application file.
+
+```python
+import agentbasis
+
+# Initialize with your API Key and Agent ID
+agentbasis.init(
+    api_key="your-api-key-here", 
+    agent_id="your-agent-id-here"
+)
+```
+
+### 2. Manual Tracking (Decorators)
+Use the `@trace` decorator to automatically track any function.
+
+```python
+from agentbasis import trace
+
+@trace
+def chat_with_user(query):
+    # Your agent logic here
+    return "Response to: " + query
+
+# When you call this, data is automatically sent to AgentBasis
+chat_with_user("Hello world")
+```
+
+### 3. OpenAI Integration
+Automatically track all your OpenAI calls (models, tokens, prompts) with one line of code.
+
+```python
+from agentbasis.llms.openai import instrument
+
+# Enable OpenAI instrumentation
+instrument()
+
+# Now just use the OpenAI client as normal
+from openai import OpenAI
+client = OpenAI()
+response = client.chat.completions.create(
+    model="gpt-4",
+    messages=[{"role": "user", "content": "Hello"}]
+)
+```
+
+### 4. Anthropic Integration
+Automatically track all your Anthropic Claude calls.
+
+```python
+from agentbasis.llms.anthropic import instrument
+
+# Enable Anthropic instrumentation
+instrument()
+
+# Now just use the Anthropic client as normal
+from anthropic import Anthropic
+client = Anthropic()
+response = client.messages.create(
+    model="claude-3-opus-20240229",
+    max_tokens=1024,
+    messages=[{"role": "user", "content": "Hello"}]
+)
+```
+
+### 5. LangChain Integration
+Track chains, tools, retrievers, and LLM calls in LangChain with full parent-child span relationships.
+
+```python
+from agentbasis.frameworks.langchain import get_callback_handler
+
+# Create a callback handler
+handler = get_callback_handler()
+
+# Pass it to your LangChain calls
+from langchain_openai import ChatOpenAI
+llm = ChatOpenAI(model="gpt-4")
+response = llm.invoke("Hello world", config={"callbacks": [handler]})
+```
+
+For chains and agents, pass the callback handler in the config:
+
+```python
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
+from agentbasis.frameworks.langchain import get_callback_config
+
+# Use get_callback_config() for convenience
+chain = LLMChain(llm=llm, prompt=PromptTemplate.from_template("{query}"))
+result = chain.invoke({"query": "What is AI?"}, config=get_callback_config())
+```
+
+### 6. Pydantic AI Integration
+Track Pydantic AI agents with built-in OpenTelemetry support.
+
+```python
+from agentbasis.frameworks.pydanticai import instrument
+
+# Enable global instrumentation for all Pydantic AI agents
+instrument()
+
+# Your agents are now automatically traced
+from pydantic_ai import Agent
+agent = Agent("openai:gpt-4")
+result = agent.run_sync("Hello!")
+```
+
+For per-agent control with user context:
+
+```python
+from agentbasis.frameworks.pydanticai import create_traced_agent
+
+# Create an agent pre-configured with tracing and context
+agent = create_traced_agent(
+    "openai:gpt-4",
+    system_prompt="You are a helpful assistant."
+)
+
+# Set user context - it will be included in traces
+agentbasis.set_user("user-123")
+result = agent.run_sync("Hello!")
+```
+
+### 7. Track Users & Sessions (Optional)
+Associate traces with specific users and sessions to debug issues and see per-user analytics.
+
+```python
+# Set the current user (from your auth system)
+agentbasis.set_user(current_user.id)
+
+# Optionally set session and conversation IDs
+agentbasis.set_session("session-abc")
+agentbasis.set_conversation("conv-123")
+
+# All subsequent LLM calls will be tagged with this context
+response = client.chat.completions.create(...)
+```
+
+Or use the context manager for scoped context:
+
+```python
+from agentbasis import context
+
+with context(user_id="user-123", session_id="session-abc"):
+    # All traces in this block include the context
+    response = client.chat.completions.create(...)
+```
+
+## Core Concepts
+
+- **OpenTelemetry**: We use OTel under the hood for maximum compatibility.
+- **Spans**: Every action (function call, LLM request) is recorded as a Span.
+- **Transport**: Data is batched and sent asynchronously to the AgentBasis backend.
+
+## Documentation
+
+For full documentation, visit [docs.agentbasis.co](https://docs.agentbasis.co).
+
+## Contributing
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup and guidelines.
+
+## License
+MIT License - see [LICENSE](LICENSE) for details.
+

agentbasis-0.1.0.dist-info/RECORD

@@ -0,0 +1,19 @@
+agentbasis/__init__.py,sha256=CskAzB1oO9YdoHkr6e4pizEtB3k8N9lMqpfv8yoslzY,2406
+agentbasis/client.py,sha256=70Z_TDNA5MjqQtz6ncJ0rRQJxyJR89qjWt9OUK_mIg8,4759
+agentbasis/config.py,sha256=PxhhcBrLA6E8UR03R2KiPqq7KV8mCdydo87i1SkaveA,1354
+agentbasis/context.py,sha256=zrHHlNa2-84eiHMrAguBXgIkLZ9xMF1TYalzf1FgsBs,7959
+agentbasis/decorators.py,sha256=jO0sWu4YaGBo0EPdgdMqM-7BPKxDw3-idWEgG-RowEo,2736
+agentbasis/frameworks/langchain/__init__.py,sha256=M0y_Ylh1Os62rVZwRkFdzwNxD_rFrybUDdvrNANzvTs,3385
+agentbasis/frameworks/langchain/callback.py,sha256=wYQ5h2BPOmIYUtqmTeS8mBoPTE9_UC0eN9KrN2N0jvU,14788
+agentbasis/frameworks/pydanticai/__init__.py,sha256=t2A0ptvypQnd_wU9fRODpnUJA8Ez3DRhTCs-kVIRkLQ,785
+agentbasis/frameworks/pydanticai/instrumentation.py,sha256=opW5OYq_FJ3zXQLdpoNdKCeX4op22R9cXhsU3Oa3HOg,7993
+agentbasis/llms/anthropic/__init__.py,sha256=kGSmvlZv9OC8-qJOzZhGfkoNveXEvdpCfBxtqIvkCUM,538
+agentbasis/llms/anthropic/messages.py,sha256=Uz1042c_bbpAcoE-9jOOH2CMQJ94x98COzWCsYC4Jno,11922
+agentbasis/llms/gemini/__init__.py,sha256=My4V_p4CeckB7JW4HIy_JSjJTia17GlgZ3HCpHgoDpA,648
+agentbasis/llms/gemini/chat.py,sha256=2J2nEvTkeuI6qXDJdM6Cq5vopmVb5wQn0Ls6cU-2-xM,13039
+agentbasis/llms/openai/__init__.py,sha256=ciehM0Qnl0HKcmnvrSvH62orE-iY_Z-l0XD56qEOCkg,594
+agentbasis/llms/openai/chat.py,sha256=HHrjuElQX4UALLJuSnu9fNJl6EssIYXiOJ5CLNpg1WQ,8773
+agentbasis-0.1.0.dist-info/METADATA,sha256=djYSYzRXAhJ8RUWUUXX96-tI2SPjEZzHZyFO9A43n5w,6721
+agentbasis-0.1.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+agentbasis-0.1.0.dist-info/top_level.txt,sha256=tqdyIQGlh-ZaQ4joP96SeZof0zSV1RwWbqrSQWl3ztc,11
+agentbasis-0.1.0.dist-info/RECORD,,

agentbasis-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

agentbasis-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+agentbasis