ragaai-catalyst 2.1b0__py3-none-any.whl → 2.1b2__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/network_tracer.py

@@ -0,0 +1,286 @@
+from datetime import datetime
+import socket
+from http.client import HTTPConnection, HTTPSConnection
+import aiohttp
+import requests
+import urllib
+import uuid
+
+
+class NetworkTracer:
+    def __init__(self):
+        self.network_calls = []
+        self.patches_applied = False  # Track whether patches are active
+        # Store original functions for restoration
+        self._original_urlopen = None
+        self._original_requests_request = None
+        self._original_http_request = None
+        self._original_https_request = None
+        self._original_socket_create_connection = None
+
+    def record_call(
+        self,
+        method,
+        url,
+        status_code=None,
+        error=None,
+        start_time=None,
+        end_time=None,
+        request_headers=None,
+        response_headers=None,
+        request_body=None,
+        response_body=None,
+    ):
+        duration = (
+            (end_time - start_time).total_seconds() if start_time and end_time else None
+        )
+        
+        # Calculate bytes sent/received from headers and body
+        bytes_sent = len(str(request_headers or "")) + len(str(request_body or ""))
+        bytes_received = len(str(response_headers or "")) + len(str(response_body or ""))
+        
+        # Extract protocol from URL
+        protocol = "https" if url.startswith("https") else "http"
+        
+        self.network_calls.append(
+            {
+                "url": url,
+                "method": method,
+                "status_code": status_code,
+                "response_time": duration,
+                "bytes_sent": bytes_sent,
+                "bytes_received": bytes_received,
+                "protocol": protocol,
+                "connection_id": str(uuid.uuid4()),  # Generate unique connection ID
+                "parent_id": None,  # Will be set by the component
+                "request": {
+                    "headers": request_headers,
+                    "body": request_body[:1000] if request_body else None,  # Limit to 1000 chars
+                },
+                "response": {
+                    "headers": response_headers,
+                    "body": response_body[:1000] if response_body else None,  # Limit to 1000 chars
+                },
+                "error": str(error) if error else None,
+            }
+        )
+
+    def activate_patches(self):
+        if not self.patches_applied:
+            # Apply monkey patches and store originals
+            self._original_urlopen = monkey_patch_urllib(self)
+            self._original_requests_request = monkey_patch_requests(self)
+            self._original_http_request, self._original_https_request = (
+                monkey_patch_http_client(self)
+            )
+            self._original_socket_create_connection = monkey_patch_socket(self)
+            self.patches_applied = True
+
+    def deactivate_patches(self):
+        if self.patches_applied:
+            # Restore original functions
+            restore_urllib(self._original_urlopen)
+            restore_requests(self._original_requests_request)
+            restore_http_client(
+                self._original_http_request, self._original_https_request
+            )
+            restore_socket(self._original_socket_create_connection)
+            self.network_calls = []
+            self.patches_applied = False
+
+
+# Define the monkey patch and restore functions
+def monkey_patch_urllib(network_tracer):
+    from urllib.request import urlopen
+
+    original_urlopen = urlopen
+
+    def patched_urlopen(url, data=None, timeout=None, *args, **kwargs):
+        if isinstance(url, str):
+            method = "GET" if data is None else "POST"
+            url_str = url
+        else:
+            method = url.get_method()
+            url_str = url.full_url
+
+        start_time = datetime.now()
+        try:
+            response = original_urlopen(url, data, timeout, *args, **kwargs)
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method=method,
+                url=url_str,
+                status_code=response.status,
+                start_time=start_time,
+                end_time=end_time,
+                request_headers=dict(response.request.headers),
+                response_headers=dict(response.headers),
+                request_body=data,
+                response_body=response.read().decode("utf-8", errors="ignore"),
+            )
+            return response
+        except Exception as e:
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method=method,
+                url=url_str,
+                error=e,
+                start_time=start_time,
+                end_time=end_time,
+            )
+            raise
+
+    urllib.request.urlopen = patched_urlopen
+    return original_urlopen  # Return the original function
+
+
+def restore_urllib(original_urlopen):
+    urllib.request.urlopen = original_urlopen
+
+
+def monkey_patch_requests(network_tracer):
+    original_request = requests.Session.request
+
+    def patched_request(self, method, url, *args, **kwargs):
+        start_time = datetime.now()
+        try:
+            response = original_request(self, method, url, *args, **kwargs)
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method=method,
+                url=url,
+                status_code=response.status_code,
+                start_time=start_time,
+                end_time=end_time,
+                request_headers=dict(response.request.headers),
+                response_headers=dict(response.headers),
+                request_body=kwargs.get("data") or kwargs.get("json"),
+                response_body=response.text,
+            )
+            return response
+        except Exception as e:
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method=method,
+                url=url,
+                error=e,
+                start_time=start_time,
+                end_time=end_time,
+            )
+            raise
+
+    requests.Session.request = patched_request
+    return original_request
+
+
+def restore_requests(original_request):
+    requests.Session.request = original_request
+
+
+def monkey_patch_http_client(network_tracer):
+    original_http_request = HTTPConnection.request
+    original_https_request = HTTPSConnection.request
+
+    def patched_request(self, method, url, body=None, headers=None, *args, **kwargs):
+        start_time = datetime.now()
+        try:
+            result = (
+                original_http_request(self, method, url, body, headers, *args, **kwargs)
+                if isinstance(self, HTTPConnection)
+                else original_https_request(
+                    self, method, url, body, headers, *args, **kwargs
+                )
+            )
+            response = self.getresponse()
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method=method,
+                url=f"{self._http_vsn_str} {self.host}:{self.port}{url}",
+                status_code=response.status,
+                start_time=start_time,
+                end_time=end_time,
+                request_headers=headers,
+                response_headers=dict(response.headers),
+                request_body=body,
+                response_body=response.read().decode("utf-8", errors="ignore"),
+            )
+            return result
+        except Exception as e:
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method=method,
+                url=f"{self._http_vsn_str} {self.host}:{self.port}{url}",
+                error=e,
+                start_time=start_time,
+                end_time=end_time,
+            )
+            raise
+
+    HTTPConnection.request = patched_request
+    HTTPSConnection.request = patched_request
+    return original_http_request, original_https_request
+
+
+def restore_http_client(original_http_request, original_https_request):
+    HTTPConnection.request = original_http_request
+    HTTPSConnection.request = original_https_request
+
+
+def monkey_patch_socket(network_tracer):
+    original_create_connection = socket.create_connection
+
+    def patched_create_connection(address, *args, **kwargs):
+        host, port = address
+        start_time = datetime.now()
+        try:
+            result = original_create_connection(address, *args, **kwargs)
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method="CONNECT",
+                url=f"{host}:{port}",
+                start_time=start_time,
+                end_time=end_time,
+            )
+            return result
+        except Exception as e:
+            end_time = datetime.now()
+            network_tracer.record_call(
+                method="CONNECT",
+                url=f"{host}:{port}",
+                error=e,
+                start_time=start_time,
+                end_time=end_time,
+            )
+            raise
+
+    socket.create_connection = patched_create_connection
+    return original_create_connection
+
+
+def restore_socket(original_create_connection):
+    socket.create_connection = original_create_connection
+
+
+async def patch_aiohttp_trace_config(network_tracer):
+    async def on_request_start(session, trace_config_ctx, params):
+        trace_config_ctx.start = datetime.now()
+
+    async def on_request_end(session, trace_config_ctx, params):
+        end_time = datetime.now()
+        response = params.response
+        network_tracer.record_call(
+            method=params.method,
+            url=str(params.url),
+            status_code=response.status,
+            start_time=trace_config_ctx.start,
+            end_time=end_time,
+            request_headers=dict(params.headers),
+            response_headers=dict(response.headers),
+            request_body=await params.response.text(),
+            response_body=await response.text(),
+        )
+
+    trace_config = aiohttp.TraceConfig()
+    trace_config.on_request_start.append(on_request_start)
+    trace_config.on_request_end.append(on_request_end)
+    return trace_config

ragaai_catalyst/tracers/agentic_tracing/sample.py

@@ -0,0 +1,197 @@
+import os
+import json
+from openai import OpenAI
+import requests
+from datetime import datetime
+from dotenv import load_dotenv
+import sys
+
+# Load environment variables
+load_dotenv()
+
+# Initialize OpenAI client
+client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+# Initialize tracer
+from tracer import Tracer
+tracer = Tracer(
+    project_name="travel_agent_demo",
+    output_dir="./traces"
+)
+
+# Start tracing
+tracer.start()
+
+@tracer.trace_tool(
+    name="llm_call",
+    tool_type="llm",
+    version="1.0.0"
+)
+def llm_call(prompt, max_tokens=512, model="gpt-3.5-turbo"):
+    response = client.chat.completions.create(
+        model=model,
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=max_tokens,
+        temperature=0.7,
+    )
+    return response.choices[0].message.content.strip()
+
+@tracer.trace_tool(
+    name="weather_tool",
+    tool_type="api",
+    version="1.0.0"
+)
+def weather_tool(destination):
+    api_key = os.environ.get("OPENWEATHERMAP_API_KEY")
+    base_url = "http://api.openweathermap.org/data/2.5/weather"
+    params = {"q": destination, "appid": api_key, "units": "metric"}
+
+    try:
+        response = requests.get(base_url, params=params)
+        response.raise_for_status()
+        data = response.json()
+        weather_description = data["weather"][0]["description"]
+        temperature = data["main"]["temp"]
+        return f"{weather_description.capitalize()}, {temperature:.1f}°C"
+    except requests.RequestException:
+        return "Weather data not available."
+
+@tracer.trace_tool(
+    name="currency_converter_tool",
+    tool_type="api",
+    version="1.0.0"
+)
+def currency_converter_tool(amount, from_currency, to_currency):
+    api_key = os.environ.get("EXCHANGERATE_API_KEY")
+    base_url = f"https://v6.exchangerate-api.com/v6/{api_key}/pair/{from_currency}/{to_currency}"
+
+    try:
+        response = requests.get(base_url)
+        response.raise_for_status()
+        data = response.json()
+        if data["result"] == "success":
+            rate = data["conversion_rate"]
+            return amount * rate
+        return None
+    except requests.RequestException:
+        return None
+
+@tracer.trace_tool(
+    name="flight_price_estimator_tool",
+    tool_type="mock",
+    version="1.0.0"
+)
+def flight_price_estimator_tool(origin, destination):
+    return f"Estimated price from {origin} to {destination}: $500-$1000"
+
+@tracer.trace_agent(
+    name="itinerary_agent",
+    agent_type="planner",
+    capabilities=["itinerary_planning", "llm_interaction"]
+)
+class ItineraryAgent:
+    def __init__(self, persona="Itinerary Agent"):
+        self.persona = persona
+
+    def plan_itinerary(self, user_preferences, duration=3):
+            itinerary_prompt = f"""
+You are a travel expert named {self.persona}.
+Based on the following user preferences, create a {duration}-day travel itinerary.
+
+User Preferences:
+{user_preferences}
+
+Itinerary:
+"""
+            return llm_call(itinerary_prompt, max_tokens=512)
+
+@tracer.trace_agent(
+    name="travel_agent",
+    agent_type="orchestrator",
+    capabilities=["preference_extraction", "travel_planning", "information_gathering"]
+)
+def travel_agent():
+    print("Welcome to the Personalized Travel Planner!\n")
+
+    # Get user input
+    user_input = "karela, 10 days, $100, nature"
+
+    # Extract preferences
+    preferences_prompt = f"""
+Extract key travel preferences from the following user input:
+"{user_input}"
+
+Please provide the extracted information in this format:
+Destination:
+Activities:
+Budget:
+Duration (in days):
+"""
+    extracted_preferences = llm_call(preferences_prompt)
+    print("\nExtracted Preferences:")
+    print(extracted_preferences)
+
+    # Parse extracted preferences
+    preferences = {}
+    for line in extracted_preferences.split("\n"):
+        if ":" in line:
+            key, value = line.split(":", 1)
+            preferences[key.strip()] = value.strip()
+
+    # Validate extracted preferences
+    required_keys = ["Destination", "Activities", "Budget", "Duration (in days)"]
+    if not all(key in preferences for key in required_keys):
+        print("\nCould not extract all required preferences. Please try again.")
+        return
+
+    # Fetch additional information
+    weather = weather_tool(preferences["Destination"])
+    print(f"\nWeather in {preferences['Destination']}: {weather}")
+
+    origin = "delhi"
+    flight_price = flight_price_estimator_tool(origin, preferences["Destination"])
+    print(flight_price)
+
+    # Plan itinerary
+    itinerary_agent = ItineraryAgent()
+    itinerary = itinerary_agent.plan_itinerary(
+        extracted_preferences, int(preferences["Duration (in days)"])
+    )
+    print("\nPlanned Itinerary:")
+    print(itinerary)
+
+    # Currency conversion
+    budget_amount = float(preferences["Budget"].replace("$", "").replace(",", ""))
+    converted_budget = currency_converter_tool(budget_amount, "USD", "INR")
+    if converted_budget:
+        print(f"\nBudget in INR: {converted_budget:.2f} INR")
+    else:
+        print("\nCurrency conversion not available.")
+
+    # Generate travel summary
+    summary_prompt = f"""
+Summarize the following travel plan:
+
+Destination: {preferences['Destination']}
+Activities: {preferences['Activities']}
+Budget: {preferences['Budget']}
+Duration: {preferences['Duration (in days)']} days
+Itinerary: {itinerary}
+Weather: {weather}
+Flight Price: {flight_price}
+
+Travel Summary:
+"""
+    travel_summary = llm_call(summary_prompt, max_tokens=2048)
+    print("\nTravel Summary:")
+    print(travel_summary)
+
+def main():
+    try:
+        travel_agent()
+    finally:
+        # Stop tracing and save results
+        tracer.stop()
+
+if __name__ == "__main__":
+    main()