add graphlit instrumentation

Author: obinnaokafor

src/.DS_Store

@@ -0,0 +1,9 @@
+import asyncio
+from examples.graphlit_example.conversation import complete
+from langtrace_python_sdk import with_langtrace_root_span
+
+
+class GraphlitRunner:
+    @with_langtrace_root_span("GraphlitRun")
+    def run(self):
+        asyncio.run(complete())

src/examples/graphlit_example/__init__.py

@@ -0,0 +1,87 @@
+import asyncio
+from dotenv import find_dotenv, load_dotenv
+from langtrace_python_sdk import langtrace
+from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+from openai import OpenAI
+from graphlit import Graphlit
+from graphlit_api import input_types, enums, exceptions
+
+_ = load_dotenv(find_dotenv())
+
+langtrace.init()
+
+
+graphlit = Graphlit()
+
+langtrace.init()
+client = OpenAI()
+
+
+async def complete():
+    uri = "https://themixchief.com"
+    try:
+        ingest_response = await graphlit.client.ingest_uri(uri=uri, is_synchronous=True)
+        content_id = ingest_response.ingest_uri.id if ingest_response.ingest_uri is not None else None
+
+        if content_id is not None:
+            print(f'Ingested content [{content_id}]:')
+
+        prompt = "In 1 sentence, what does mixchief do."
+
+        model = "gpt-4o"
+
+        specification_input = input_types.SpecificationInput(
+            name=f"OpenAI [{str(enums.OpenAIModels.GPT4O_128K)}]",
+            type=enums.SpecificationTypes.COMPLETION,
+            serviceType=enums.ModelServiceTypes.OPEN_AI,
+            openAI=input_types.OpenAIModelPropertiesInput(
+                model=enums.OpenAIModels.GPT4O_128K,
+            )
+        )
+
+        specification_response = await graphlit.client.create_specification(specification_input)
+        specification_id = specification_response.create_specification.id if specification_response.create_specification is not None else None
+
+        if specification_id is not None:
+            print(f'Created specification [{specification_id}].')
+
+            conversation_input = input_types.ConversationInput(
+                name="Conversation",
+                specification=input_types.EntityReferenceInput(
+                    id=specification_id
+                ),
+            )
+
+            conversation_response = await graphlit.client.create_conversation(conversation_input)
+            conversation_id = conversation_response.create_conversation.id if conversation_response.create_conversation is not None else None
+
+            if conversation_id is not None:
+                print(f'Created conversation [{conversation_id}].')
+
+                format_response = await graphlit.client.format_conversation(prompt, conversation_id)
+                formatted_message = format_response.format_conversation.message.message if format_response.format_conversation is not None and format_response.format_conversation.message is not None else None
+
+                if formatted_message is not None:
+                    stream_response = client.chat.completions.create(
+                        model=model,
+                        messages=[{"role": "user", "content": formatted_message}],
+                        temperature=0.1,
+                        top_p=0.2,
+                        stream=True
+                    )
+
+                    completion = ""
+
+                    for chunk in stream_response:
+                        delta = chunk.choices[0].delta.content
+
+                        if delta is not None:
+                            completion += delta
+
+                    if completion is not None:
+                        # NOTE: stores completion back into conversation
+                        complete_response = await graphlit.client.complete_conversation(completion, conversation_id)
+
+                        print(complete_response.complete_conversation.message.message if complete_response.complete_conversation is not None and complete_response.complete_conversation.message is not None else "None")
+    except exceptions.GraphQLClientError as e:
+        print(f"Graphlit API error: {e}")

src/examples/graphlit_example/conversation.py

@@ -40,6 +40,7 @@
     "AWS_BEDROCK": "AWS Bedrock",
     "CEREBRAS": "Cerebras",
     "MILVUS": "Milvus",
+    "GRAPHLIT": "Graphlit",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -26,6 +26,10 @@
 from .qdrant import QdrantInstrumentation
 from .vertexai import VertexAIInstrumentation
 from .weaviate import WeaviateInstrumentation
+from .cerebras import CerebrasInstrumentation
+from .milvus import MilvusInstrumentation
+from .google_genai import GoogleGenaiInstrumentation
+from .graphlit import GraphlitInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -56,4 +60,5 @@
     "MilvusInstrumentation",
     "GoogleGenaiInstrumentation",
     "CrewaiToolsInstrumentation",
+    "GraphlitInstrumentation",
 ]

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import GraphlitInstrumentation
+
+__all__ = ["GraphlitInstrumentation"]

src/langtrace_python_sdk/instrumentation/graphlit/__init__.py

@@ -0,0 +1,57 @@
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper as _W
+from typing import Collection
+from importlib_metadata import version as v
+from .patch import patch_graphlit_operation
+
+class GraphlitInstrumentation(BaseInstrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["graphlit-client >= 1.0.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = v("graphlit-client")
+        try:
+            _W(
+                "graphlit.graphlit",
+                "Client.ingest_uri",
+                patch_graphlit_operation("ingest_uri", version, tracer),
+            )
+            _W(
+                "graphlit.graphlit",
+                "Client.create_feed",
+                patch_graphlit_operation("create_feed", version, tracer),
+            )
+            _W(
+                "graphlit.graphlit",
+                "Client.create_specification",
+                patch_graphlit_operation("create_specification", version, tracer),
+            )
+            _W(
+                "graphlit.graphlit",
+                "Client.create_conversation",
+                patch_graphlit_operation("create_conversation", version, tracer),
+            )
+            _W(
+                "graphlit.graphlit",
+                "Client.format_conversation",
+                patch_graphlit_operation("format_conversation", version, tracer),
+            )
+            _W(
+                "graphlit.graphlit",
+                "Client.complete_conversation",
+                patch_graphlit_operation("complete_conversation", version, tracer),
+            )
+            _W(
+                "graphlit.graphlit",
+                "Client.prompt_conversation",
+                patch_graphlit_operation("prompt_conversation", version, tracer),
+            )
+        except Exception:
+            pass
+
+    def _uninstrument(self, **kwargs):
+        pass

src/langtrace_python_sdk/instrumentation/graphlit/instrumentation.py

@@ -0,0 +1,69 @@
+import json
+from importlib_metadata import version as v
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry import baggage
+from opentelemetry.trace import Span, SpanKind, Tracer
+from opentelemetry.trace.status import Status, StatusCode
+
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
+from langtrace_python_sdk.utils.llm import set_span_attributes
+from langtrace_python_sdk.utils.misc import serialize_args, serialize_kwargs
+
+
+def patch_graphlit_operation(operation_name, version, tracer: Tracer):
+    async def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["GRAPHLIT"]
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "framework",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        span_attributes["langchain.metadata"] = serialize_kwargs(**kwargs)
+        span_attributes["langchain.inputs"] = serialize_args(*args)
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=f"graphlit.{operation_name}", kind=SpanKind.CLIENT
+        ) as span:
+            try:
+                set_span_attributes(span, attributes)
+                result = await wrapped(*args, **kwargs)
+                
+                if result:
+                    operation_result = json.loads(result.model_dump_json())[operation_name]
+                    if operation_name == "complete_conversation" or operation_name == "prompt_conversation" or operation_name == "format_conversation":
+                        set_graphlit_conversation_attributes(span, operation_result)
+                    else:
+                        for key, value in operation_result.items():
+                            span.set_attribute(f"graphlit.{operation_name}.{key}", str(value))
+
+                    span.set_status(Status(StatusCode.OK))
+
+                return result
+
+            except Exception as err:
+                span.record_exception(err)
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+                raise
+
+    return traced_method
+
+
+def set_graphlit_conversation_attributes(span: Span, response):
+    if not response or "message" not in response:
+        return
+
+    span.set_attribute(f"graphlit.conversation.id", response['conversation']['id'])
+
+    for key, value in response['message'].items():
+        span.set_attribute(f"graphlit.conversation.{key}", str(value))

src/langtrace_python_sdk/instrumentation/graphlit/patch.py

@@ -29,9 +29,58 @@
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
 from opentelemetry.sdk.trace import TracerProvider
+<<<<<<< HEAD
 from opentelemetry.sdk.trace.export import (BatchSpanProcessor,
                                             ConsoleSpanExporter,
                                             SimpleSpanProcessor)
+=======
+from opentelemetry.sdk.trace.export import (
+    BatchSpanProcessor,
+    ConsoleSpanExporter,
+    SimpleSpanProcessor,
+)
+
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+    OTLPSpanExporter as GRPCExporter,
+)
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+    OTLPSpanExporter as HTTPExporter,
+)
+from langtrace_python_sdk.constants.exporter.langtrace_exporter import (
+    LANGTRACE_REMOTE_URL,
+    LANGTRACE_SESSION_ID_HEADER,
+)
+from langtrace_python_sdk.instrumentation import (
+    AnthropicInstrumentation,
+    ChromaInstrumentation,
+    CohereInstrumentation,
+    CrewAIInstrumentation,
+    DspyInstrumentation,
+    EmbedchainInstrumentation,
+    GeminiInstrumentation,
+    GroqInstrumentation,
+    LangchainCommunityInstrumentation,
+    LangchainCoreInstrumentation,
+    LangchainInstrumentation,
+    LanggraphInstrumentation,
+    LiteLLMInstrumentation,
+    LlamaindexInstrumentation,
+    MistralInstrumentation,
+    AWSBedrockInstrumentation,
+    OllamaInstrumentor,
+    OpenAIInstrumentation,
+    PineconeInstrumentation,
+    QdrantInstrumentation,
+    AutogenInstrumentation,
+    VertexAIInstrumentation,
+    WeaviateInstrumentation,
+    PyMongoInstrumentation,
+    CerebrasInstrumentation,
+    MilvusInstrumentation,
+    GoogleGenaiInstrumentation,
+    GraphlitInstrumentation,
+)
+>>>>>>> 0f7e0f0 (add graphlit instrumentation)
 from opentelemetry.util.re import parse_env_headers
 from sentry_sdk.types import Event, Hint
 
@@ -279,6 +328,7 @@ def init(
         "google-cloud-aiplatform": VertexAIInstrumentation(),
         "google-generativeai": GeminiInstrumentation(),
         "google-genai": GoogleGenaiInstrumentation(),
+        "graphlit-client": GraphlitInstrumentation(),
         "mistralai": MistralInstrumentation(),
         "boto3": AWSBedrockInstrumentation(),
         "autogen": AutogenInstrumentation(),

src/langtrace_python_sdk/langtrace.py

@@ -20,9 +20,10 @@
     "vertexai": False,
     "gemini": False,
     "mistral": False,
-    "awsbedrock": True,
+    "awsbedrock": False,
     "cerebras": False,
     "google_genai": False,
+    "graphlit": True,
 }
 
 if ENABLED_EXAMPLES["anthropic"]:
@@ -144,3 +145,9 @@
 
     print(Fore.BLUE + "Running Google Genai example" + Fore.RESET)
     GoogleGenaiRunner().run()
+
+if ENABLED_EXAMPLES["graphlit"]:
+    from examples.graphlit_example import GraphlitRunner
+
+    print(Fore.BLUE + "Running Graphlit example" + Fore.RESET)
+    GraphlitRunner().run()