Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: added OTLP export #38

Merged
merged 3 commits into from
Jan 11, 2024
+537 −19
Merged

fix: added OTLP export #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a87e3b2
fix: added otlp export
adubovik Jan 11, 2024
fa7eb03
fix: added logging for tracing
adubovik Jan 11, 2024
40b755b
fix: added httpx as dependency
adubovik Jan 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 25 additions & 12 deletions aidial_adapter_vertexai/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import json
import logging.config
import os
from typing import Optional

from aidial_sdk import DIALApp
from aidial_sdk import HTTPException as DialException
from fastapi import Body, Header, Path, Request, Response
from aidial_sdk.telemetry.types import TelemetryConfig, TracingConfig
from fastapi import Body, Header, Path, Request
from fastapi.responses import JSONResponse

from aidial_adapter_vertexai.chat_completion import VertexAIChatCompletion
Expand All @@ -29,15 +31,26 @@

logging.config.dictConfig(LogConfig().dict())

region = get_env("DEFAULT_REGION")
gcp_project_id = get_env("GCP_PROJECT_ID")

app = DIALApp(description="Google VertexAI adapter for DIAL API")
DEFAULT_REGION = get_env("DEFAULT_REGION")
GCP_PROJECT_ID = get_env("GCP_PROJECT_ID")

OTLP_EXPORT_ENABLED: bool = (
os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") is not None
)

@app.get("/healthcheck")
def healthcheck():
return Response("OK")
app = DIALApp(
description="Google VertexAI adapter for DIAL API",
add_healthcheck=True,
telemetry_config=TelemetryConfig(
service_name="bedrock",
tracing=TracingConfig(
otlp_export=OTLP_EXPORT_ENABLED,
logging=True,
),
)
if OTLP_EXPORT_ENABLED
else None,
)


@app.get("/openai/models")
Expand All @@ -55,8 +68,8 @@ async def models():
app.add_chat_completion(
deployment.get_model_id(),
VertexAIChatCompletion(
project_id=gcp_project_id,
region=region,
project_id=GCP_PROJECT_ID,
region=DEFAULT_REGION,
),
)

Expand All @@ -76,9 +89,9 @@ async def embeddings(
log.debug(f"query:\n{json.dumps(query.dict(exclude_none=True))}")

model = await get_embeddings_model(
location=region,
location=DEFAULT_REGION,
deployment=deployment,
project_id=gcp_project_id,
project_id=GCP_PROJECT_ID,
)

response = await model.embeddings(
Expand Down
6 changes: 4 additions & 2 deletions aidial_adapter_vertexai/llm/gemini_prompt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import List
from typing import List, assert_never

from aidial_sdk.chat_completion import Message, Role
from pydantic import BaseModel
Expand Down Expand Up @@ -54,8 +54,10 @@ def from_message(cls, message: Message) -> "SimpleMessage":
role = ChatSession._USER_ROLE
case Role.ASSISTANT:
role = ChatSession._MODEL_ROLE
case Role.FUNCTION:
case Role.FUNCTION | Role.TOOL:
raise ValidationError("Function messages are not supported")
case _:
assert_never(message.role)

return SimpleMessage(role=role, content=content)

Expand Down
Loading