Integrate AgentOps with ReAG

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

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"

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
+    )

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]:

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