merge

karthikscale3 · karthikscale3 · commit fd9baed3d698 · 2025-02-06T22:35:18.000+05:30
diff --git a/pyproject.toml b/pyproject.toml
@@ -18,7 +18,7 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 dependencies = [
-  'trace-attributes==7.1.1',
+  'trace-attributes==7.2.0',
   'opentelemetry-api>=1.25.0',
   'opentelemetry-sdk>=1.25.0',
   'opentelemetry-instrumentation>=0.47b0',
@@ -44,6 +44,7 @@ dev = [
   "anthropic",
   "chromadb",
   "qdrant-client",
+  "graphlit-client",
   "python-dotenv",
   "pinecone-client",
   "langchain",
diff --git a/src/.DS_Store b/src/.DS_Store
diff --git a/src/examples/graphlit_example/__init__.py b/src/examples/graphlit_example/__init__.py
@@ -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())
diff --git a/src/examples/graphlit_example/conversation.py b/src/examples/graphlit_example/conversation.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}")
diff --git a/src/langtrace_python_sdk/constants/instrumentation/common.py b/src/langtrace_python_sdk/constants/instrumentation/common.py
@@ -30,7 +30,7 @@
     "QDRANT": "Qdrant",
     "WEAVIATE": "Weaviate",
     "OLLAMA": "Ollama",
-    "VERTEXAI": "VertexAI",
+    "VERTEXAI": "Vertex AI",
     "GEMINI": "Gemini",
     "MISTRAL": "Mistral",
     "EMBEDCHAIN": "Embedchain",
@@ -40,6 +40,7 @@
     "AWS_BEDROCK": "AWS Bedrock",
     "CEREBRAS": "Cerebras",
     "MILVUS": "Milvus",
+    "GRAPHLIT": "Graphlit",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"
diff --git a/src/langtrace_python_sdk/constants/instrumentation/vertexai.py b/src/langtrace_python_sdk/constants/instrumentation/vertexai.py
@@ -39,4 +39,24 @@
         "method": "ChatSession",
         "operation": "send_message_streaming",
     },
+    "PREDICTION_SERVICE_BETA_GENERATE_CONTENT": {
+        "module": "google.cloud.aiplatform_v1beta1.services.prediction_service.client",
+        "method": "PredictionServiceClient",
+        "operation": "generate_content",
+    },
+    "PREDICTION_SERVICE_GENERATE_CONTENT": {
+        "module": "google.cloud.aiplatform_v1.services.prediction_service.client",
+        "method": "PredictionServiceClient",
+        "operation": "generate_content",
+    },
+    "PREDICTION_SERVICE_BETA_STREAM_GENERATE_CONTENT": {
+        "module": "google.cloud.aiplatform_v1beta1.services.prediction_service.client",
+        "method": "PredictionServiceClient",
+        "operation": "stream_generate_content",
+    },
+    "PREDICTION_SERVICE_STREAM_GENERATE_CONTENT": {
+        "module": "google.cloud.aiplatform_v1.services.prediction_service.client",
+        "method": "PredictionServiceClient",
+        "operation": "stream_generate_content",
+    },
 }
diff --git a/src/langtrace_python_sdk/instrumentation/__init__.py b/src/langtrace_python_sdk/instrumentation/__init__.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",
 ]
diff --git a/src/langtrace_python_sdk/instrumentation/graphlit/__init__.py b/src/langtrace_python_sdk/instrumentation/graphlit/__init__.py
@@ -0,0 +1,3 @@
+from .instrumentation import GraphlitInstrumentation
+
+__all__ = ["GraphlitInstrumentation"]
diff --git a/src/langtrace_python_sdk/instrumentation/graphlit/instrumentation.py b/src/langtrace_python_sdk/instrumentation/graphlit/instrumentation.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
diff --git a/src/langtrace_python_sdk/instrumentation/graphlit/patch.py b/src/langtrace_python_sdk/instrumentation/graphlit/patch.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))
diff --git a/src/langtrace_python_sdk/instrumentation/vertexai/patch.py b/src/langtrace_python_sdk/instrumentation/vertexai/patch.py
@@ -27,12 +27,13 @@ def patch_vertexai(name, version, tracer: Tracer):
     def traced_method(wrapped, instance, args, kwargs):
         service_provider = SERVICE_PROVIDERS["VERTEXAI"]
         prompts = serialize_prompts(args, kwargs)
+
         span_attributes = {
             **get_langtrace_attributes(version, service_provider),
             **get_llm_request_attributes(
                 kwargs,
                 prompts=prompts,
-                model=get_llm_model(instance),
+                model=get_llm_model(instance, kwargs),
             ),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: "",
@@ -77,6 +78,10 @@ def set_response_attributes(span: Span, result):
     if hasattr(result, "text"):
         set_event_completion(span, [{"role": "assistant", "content": result.text}])
 
+    if hasattr(result, "candidates"):
+        parts = result.candidates[0].content.parts
+        set_event_completion(span, [{"role": "assistant", "content": parts[0].text}])
+
     if hasattr(result, "usage_metadata") and result.usage_metadata is not None:
         usage = result.usage_metadata
         input_tokens = usage.prompt_token_count
@@ -96,17 +101,23 @@ def set_response_attributes(span: Span, result):
 
 
 def is_streaming_response(response):
-    return isinstance(response, types.GeneratorType) or isinstance(
-        response, types.AsyncGeneratorType
+    return (
+        isinstance(response, types.GeneratorType)
+        or isinstance(response, types.AsyncGeneratorType)
+        or str(type(response).__name__) == "_StreamingResponseIterator"
     )
 
 
-def get_llm_model(instance):
+def get_llm_model(instance, kwargs):
+    if "request" in kwargs:
+        return kwargs.get("request").model.split("/")[-1]
+
     if hasattr(instance, "_model_name"):
         return instance._model_name.replace("publishers/google/models/", "")
     return getattr(instance, "_model_id", "unknown")
 
 
+@silently_fail
 def serialize_prompts(args, kwargs):
     if args and len(args) > 0:
         prompt_parts = []
@@ -122,5 +133,24 @@ def serialize_prompts(args, kwargs):
 
         return [{"role": "user", "content": "\n".join(prompt_parts)}]
     else:
-        content = kwargs.get("prompt") or kwargs.get("message")
-        return [{"role": "user", "content": content}] if content else []
+        # Handle PredictionServiceClient for google-cloud-aiplatform.
+        if "request" in kwargs:
+            prompt = []
+            prompt_body = kwargs.get("request")
+            if prompt_body.system_instruction:
+                for part in prompt_body.system_instruction.parts:
+                    prompt.append({"role": "system", "content": part.text})
+
+            contents = prompt_body.contents
+
+            if not contents:
+                return []
+
+            for c in contents:
+                role = c.role if c.role else "user"
+                content = c.parts[0].text if c.parts else ""
+                prompt.append({"role": role, "content": content})
+            return prompt
+        else:
+            content = kwargs.get("prompt") or kwargs.get("message")
+            return [{"role": "user", "content": content}] if content else []
diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py
diff --git a/src/langtrace_python_sdk/utils/llm.py b/src/langtrace_python_sdk/utils/llm.py
diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py
diff --git a/src/run_example.py b/src/run_example.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+from .instrumentation import GraphlitInstrumentation`
	`2`	`+`
	`3`	`+__all__ = ["GraphlitInstrumentation"]`