diff --git a/python/README.md b/python/README.md index e8e09da..284dcd3 100644 --- a/python/README.md +++ b/python/README.md @@ -9,6 +9,13 @@ poetry add reag ``` +3. For AgentOps integration, install with the optional dependency: + ```bash + pip install reag[agentops] + # or + poetry add reag[agentops] + ``` + ## Quick Start ```python from reag.client import ReagClient, Document @@ -31,6 +38,65 @@ async with ReagClient( ``` +## AgentOps Integration + +ReAG supports AgentOps for monitoring and tracing your agent workflows. To enable AgentOps integration: + +1. Install ReAG with AgentOps support: + ```bash + pip install reag[agentops] + ``` + +2. Initialize AgentOps before using ReAG: + ```python + import reag + + # Initialize AgentOps + reag.init_agentops( + api_key="your-agentops-api-key", # or set AGENTOPS_API_KEY env var + tags=["reag", "rag", "llm"], + trace_name="My ReAG Workflow" + ) + + # Use ReAG as normal - it will automatically track LLM calls and operations + from reag.client import ReagClient, Document + + async with ReagClient(model="gpt-4o-mini") as client: + docs = [Document(name="test", content="test content")] + response = await client.query("What is this about?", documents=docs) + ``` + +3. For advanced usage with manual trace management: + ```python + import reag + import agentops + + # Initialize without auto-starting session + reag.init_agentops( + api_key="your-agentops-api-key", + auto_start_session=False + ) + + # Manually manage traces + tracer = agentops.start_trace(trace_name="Custom Workflow", tags=["custom"]) + + try: + async with ReagClient(model="gpt-4o-mini") as client: + docs = [Document(name="test", content="test content")] + response = await client.query("What is this about?", documents=docs) + + agentops.end_trace(tracer, end_state="Success") + except Exception as e: + agentops.end_trace(tracer, end_state="Fail") + raise + ``` + +The AgentOps integration automatically tracks: +- LLM calls made through LiteLLM +- Document filtering operations +- Query processing operations +- Content extraction operations + ## API Reference ### Initialization diff --git a/python/pyproject.toml b/python/pyproject.toml index aafd3ac..9dc0954 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -29,6 +29,10 @@ python = ">=3.9,<3.13" pydantic = "^2.0.0" httpx = "^0.25.0" litellm = "^1.60.0" +agentops = {version = "^0.3.0", optional = true} + +[tool.poetry.extras] +agentops = ["agentops"] [tool.poetry.group.dev.dependencies] pytest = "^7.4.0" diff --git a/python/src/reag/__init__.py b/python/src/reag/__init__.py index e69de29..8cc17c8 100644 --- a/python/src/reag/__init__.py +++ b/python/src/reag/__init__.py @@ -0,0 +1,47 @@ +import os +from typing import Optional + +try: + import agentops + AGENTOPS_AVAILABLE = True +except ImportError: + AGENTOPS_AVAILABLE = False + +from .client import ReagClient, Document, QueryResult, MetadataFilter +from .schema import ResponseSchema, ResponseSchemaMessage + +__version__ = "0.0.6" +__all__ = ["ReagClient", "Document", "QueryResult", "MetadataFilter", "ResponseSchema", "ResponseSchemaMessage", "init_agentops"] + +def init_agentops( + api_key: Optional[str] = None, + tags: Optional[list] = None, + trace_name: Optional[str] = None, + auto_start_session: bool = True +): + """ + Initialize AgentOps for monitoring and tracing. + + Args: + api_key: AgentOps API key. If not provided, will look for AGENTOPS_API_KEY environment variable. + tags: List of tags to apply to traces. + trace_name: Name for the trace. + auto_start_session: Whether to automatically start a session. + """ + if not AGENTOPS_AVAILABLE: + raise ImportError("AgentOps is not installed. Install it with: pip install agentops") + + # Get API key from parameter or environment + if api_key is None: + api_key = os.getenv("AGENTOPS_API_KEY") + + if api_key is None: + raise ValueError("AgentOps API key is required. Pass it as a parameter or set AGENTOPS_API_KEY environment variable.") + + # Initialize AgentOps + agentops.init( + api_key=api_key, + tags=tags or ["reag"], + trace_name=trace_name or "ReAG Workflow", + auto_start_session=auto_start_session + ) \ No newline at end of file diff --git a/python/src/reag/client.py b/python/src/reag/client.py index b7d4a3e..aad04d6 100644 --- a/python/src/reag/client.py +++ b/python/src/reag/client.py @@ -4,6 +4,29 @@ import re from typing import List, Optional, TypeVar, Dict, Union from pydantic import BaseModel + +try: + import agentops + from agentops.sdk.decorators import agent, operation, tool + AGENTOPS_AVAILABLE = True +except ImportError: + AGENTOPS_AVAILABLE = False + + def agent(*args, **kwargs): + def decorator(cls): + return cls + return decorator + + def operation(*args, **kwargs): + def decorator(func): + return func + return decorator + + def tool(*args, **kwargs): + def decorator(func): + return func + return decorator + from litellm import acompletion from reag.prompt import REAG_SYSTEM_PROMPT @@ -35,6 +58,7 @@ class QueryResult(BaseModel): DEFAULT_BATCH_SIZE = 20 +@agent(name="ReagClient") class ReagClient: def __init__( self, @@ -59,6 +83,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): if self._http_client: await self._http_client.aclose() + @operation def _filter_documents_by_metadata( self, documents: List[Document], filters: Optional[List[MetadataFilter]] = None ) -> List[Document]: @@ -129,6 +154,7 @@ def _filter_documents_by_metadata( return filtered_docs + @operation def _extract_think_content(self, text: str) -> tuple[str, str, bool]: """Extract content from think tags and parse the bulleted response format.""" # Extract think content @@ -154,6 +180,7 @@ def _extract_think_content(self, text: str) -> tuple[str, str, bool]: return content, reasoning, is_irrelevant + @operation async def query( self, prompt: str, documents: List[Document], options: Optional[Dict] = None ) -> List[QueryResult]: diff --git a/python/tests/test_client.py b/python/tests/test_client.py index de91999..e06b55b 100644 --- a/python/tests/test_client.py +++ b/python/tests/test_client.py @@ -110,3 +110,44 @@ async def test_query_returns_empty_for_irrelevant_docs(): print(response) assert response is not None assert len(response) == 0 # Should be empty since doc is irrelevant + + +@pytest.mark.asyncio +async def test_agentops_integration(): + """Test that AgentOps integration works when available""" + try: + import reag + + # This should not raise an error even if AgentOps is not installed + # due to the fallback decorators + async with ReagClient() as client: + docs = [ + Document( + name="Test Doc", + content="This is a test document for AgentOps integration.", + metadata={"type": "test"}, + ), + ] + response = await client.query("What is this document about?", documents=docs) + assert response is not None + except ImportError: + # If AgentOps is not installed, the test should still pass + # due to the fallback decorators + pass + + +def test_init_agentops_function(): + """Test that init_agentops function works correctly""" + import reag + + # Test that the function exists and is callable + assert hasattr(reag, 'init_agentops') + assert callable(reag.init_agentops) + + # Test that it raises error when AgentOps is not available and no api_key + try: + reag.init_agentops() + # Should either work (if AgentOps is installed) or raise ImportError + except (ImportError, ValueError): + # Expected if AgentOps is not installed or no API key provided + pass