Skip to content

Commit caf330e

Browse files
alizenhomalizenhom
committed
add support for langchain-vertexai
1 parent c8bca2f commit caf330e

File tree

3 files changed

+86
-19
lines changed

3 files changed

+86
-19
lines changed

‎src/langtrace_python_sdk/constants/instrumentation/vertexai.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,24 @@
3939
"method": "ChatSession",
4040
"operation": "send_message_streaming",
4141
},
42+
"PREDICTION_SERVICE_BETA_GENERATE_CONTENT": {
43+
"module": "google.cloud.aiplatform_v1beta1.services.prediction_service.client",
44+
"method": "PredictionServiceClient",
45+
"operation": "generate_content",
46+
},
47+
"PREDICTION_SERVICE_GENERATE_CONTENT": {
48+
"module": "google.cloud.aiplatform_v1.services.prediction_service.client",
49+
"method": "PredictionServiceClient",
50+
"operation": "generate_content",
51+
},
52+
"PREDICTION_SERVICE_BETA_STREAM_GENERATE_CONTENT": {
53+
"module": "google.cloud.aiplatform_v1beta1.services.prediction_service.client",
54+
"method": "PredictionServiceClient",
55+
"operation": "stream_generate_content",
56+
},
57+
"PREDICTION_SERVICE_STREAM_GENERATE_CONTENT": {
58+
"module": "google.cloud.aiplatform_v1.services.prediction_service.client",
59+
"method": "PredictionServiceClient",
60+
"operation": "stream_generate_content",
61+
},
4262
}

‎src/langtrace_python_sdk/instrumentation/vertexai/patch.py‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,13 @@ def patch_vertexai(name, version, tracer: Tracer):
2727
def traced_method(wrapped, instance, args, kwargs):
2828
service_provider = SERVICE_PROVIDERS["VERTEXAI"]
2929
prompts = serialize_prompts(args, kwargs)
30+
3031
span_attributes = {
3132
**get_langtrace_attributes(version, service_provider),
3233
**get_llm_request_attributes(
3334
kwargs,
3435
prompts=prompts,
35-
model=get_llm_model(instance),
36+
model=get_llm_model(instance, kwargs),
3637
),
3738
**get_llm_url(instance),
3839
SpanAttributes.LLM_PATH: "",
@@ -77,6 +78,10 @@ def set_response_attributes(span: Span, result):
7778
if hasattr(result, "text"):
7879
set_event_completion(span, [{"role": "assistant", "content": result.text}])
7980

81+
if hasattr(result, "candidates"):
82+
parts = result.candidates[0].content.parts
83+
set_event_completion(span, [{"role": "assistant", "content": parts[0].text}])
84+
8085
if hasattr(result, "usage_metadata") and result.usage_metadata is not None:
8186
usage = result.usage_metadata
8287
input_tokens = usage.prompt_token_count
@@ -96,17 +101,23 @@ def set_response_attributes(span: Span, result):
96101

97102

98103
def is_streaming_response(response):
99-
return isinstance(response, types.GeneratorType) or isinstance(
100-
response, types.AsyncGeneratorType
104+
return (
105+
isinstance(response, types.GeneratorType)
106+
or isinstance(response, types.AsyncGeneratorType)
107+
or str(type(response).__name__) == "_StreamingResponseIterator"
101108
)
102109

103110

104-
def get_llm_model(instance):
111+
def get_llm_model(instance, kwargs):
112+
if "request" in kwargs:
113+
return kwargs.get("request").model.split("/")[-1]
114+
105115
if hasattr(instance, "_model_name"):
106116
return instance._model_name.replace("publishers/google/models/", "")
107117
return getattr(instance, "_model_id", "unknown")
108118

109119

120+
@silently_fail
110121
def serialize_prompts(args, kwargs):
111122
if args and len(args) > 0:
112123
prompt_parts = []
@@ -122,5 +133,24 @@ def serialize_prompts(args, kwargs):
122133

123134
return [{"role": "user", "content": "\n".join(prompt_parts)}]
124135
else:
125-
content = kwargs.get("prompt") or kwargs.get("message")
126-
return [{"role": "user", "content": content}] if content else []
136+
# Handle PredictionServiceClient for google-cloud-aiplatform.
137+
if "request" in kwargs:
138+
prompt = []
139+
prompt_body = kwargs.get("request")
140+
if prompt_body.system_instruction:
141+
for part in prompt_body.system_instruction.parts:
142+
prompt.append({"role": "system", "content": part.text})
143+
144+
contents = prompt_body.contents
145+
146+
if not contents:
147+
return []
148+
149+
for c in contents:
150+
role = c.role if c.role else "user"
151+
content = c.parts[0].text if c.parts else ""
152+
prompt.append({"role": role, "content": content})
153+
return prompt
154+
else:
155+
content = kwargs.get("prompt") or kwargs.get("message")
156+
return [{"role": "user", "content": content}] if content else []

‎src/langtrace_python_sdk/utils/llm.py‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -395,18 +395,30 @@ def build_streaming_response(self, chunk):
395395
content = [chunk.text]
396396

397397
# CohereV2
398-
if (hasattr(chunk, "delta") and
399-
chunk.delta is not None and
400-
hasattr(chunk.delta, "message") and
401-
chunk.delta.message is not None and
402-
hasattr(chunk.delta.message, "content") and
403-
chunk.delta.message.content is not None and
404-
hasattr(chunk.delta.message.content, "text") and
405-
chunk.delta.message.content.text is not None):
398+
if (
399+
hasattr(chunk, "delta")
400+
and chunk.delta is not None
401+
and hasattr(chunk.delta, "message")
402+
and chunk.delta.message is not None
403+
and hasattr(chunk.delta.message, "content")
404+
and chunk.delta.message.content is not None
405+
and hasattr(chunk.delta.message.content, "text")
406+
and chunk.delta.message.content.text is not None
407+
):
406408
content = [chunk.delta.message.content.text]
407-
409+
# google-cloud-aiplatform
410+
if hasattr(chunk, "candidates") and chunk.candidates is not None:
411+
for candidate in chunk.candidates:
412+
if hasattr(candidate, "content") and candidate.content is not None:
413+
for part in candidate.content.parts:
414+
if hasattr(part, "text") and part.text is not None:
415+
content.append(part.text)
408416
# Anthropic
409-
if hasattr(chunk, "delta") and chunk.delta is not None and not hasattr(chunk.delta, "message"):
417+
if (
418+
hasattr(chunk, "delta")
419+
and chunk.delta is not None
420+
and not hasattr(chunk.delta, "message")
421+
):
410422
content = [chunk.delta.text] if hasattr(chunk.delta, "text") else []
411423

412424
if isinstance(chunk, dict):
@@ -425,9 +437,14 @@ def set_usage_attributes(self, chunk):
425437

426438
# CohereV2
427439
if hasattr(chunk, "type") and chunk.type == "message-end":
428-
if (hasattr(chunk, "delta") and chunk.delta is not None and
429-
hasattr(chunk.delta, "usage") and chunk.delta.usage is not None and
430-
hasattr(chunk.delta.usage, "billed_units") and chunk.delta.usage.billed_units is not None):
440+
if (
441+
hasattr(chunk, "delta")
442+
and chunk.delta is not None
443+
and hasattr(chunk.delta, "usage")
444+
and chunk.delta.usage is not None
445+
and hasattr(chunk.delta.usage, "billed_units")
446+
and chunk.delta.usage.billed_units is not None
447+
):
431448
usage = chunk.delta.usage.billed_units
432449
self.completion_tokens = int(usage.output_tokens)
433450
self.prompt_tokens = int(usage.input_tokens)

0 commit comments

Comments
 (0)