fix(galileo): support hosted v2 spans API and string output extraction

litellm/integrations/galileo.py

@@ -1,18 +1,24 @@
+import json
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from pydantic import BaseModel, Field
 
 import litellm
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    convert_content_list_to_str,
+    get_content_from_model_response,
+)
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
 
+GALILEO_CLOUD_API_BASE_URL = "https://api.galileo.ai"
+
 
-# from here: https://docs.rungalileo.io/galileo/gen-ai-studio-products/galileo-observe/how-to/logging-data-via-restful-apis#structuring-your-records
 class LLMResponse(BaseModel):
     latency_ms: int
     status_code: int
@@ -37,65 +43,184 @@ class GalileoObserve(CustomLogger):
     def __init__(self) -> None:
         self.in_memory_records: List[dict] = []
         self.batch_size = 1
-        self.base_url = os.getenv("GALILEO_BASE_URL", None)
-        self.project_id = os.getenv("GALILEO_PROJECT_ID", None)
+        self.api_key = os.getenv("GALILEO_API_KEY")
+        self.project_id = os.getenv("GALILEO_PROJECT_ID")
+        self.log_stream_id = os.getenv("GALILEO_LOG_STREAM_ID")
+        self.username = os.getenv("GALILEO_USERNAME")
+        self.password = os.getenv("GALILEO_PASSWORD")
+        self.base_url = self._normalize_base_url(os.getenv("GALILEO_BASE_URL"))
+        if self.api_key and not self.base_url:
+            self.base_url = GALILEO_CLOUD_API_BASE_URL
+        self.use_v2_api = bool(self.api_key)
         self.headers: Optional[Dict[str, str]] = None
         self.async_httpx_handler = get_async_httpx_client(
             llm_provider=httpxSpecialProvider.LoggingCallback
         )
-        pass
 
-    def set_galileo_headers(self):
-        # following https://docs.rungalileo.io/galileo/gen-ai-studio-products/galileo-observe/how-to/logging-data-via-restful-apis#logging-your-records
+    @staticmethod
+    def _normalize_base_url(base_url: Optional[str]) -> Optional[str]:
+        if base_url:
+            return base_url.rstrip("/")
+        return None
 
-        headers = {
-            "accept": "application/json",
-            "Content-Type": "application/x-www-form-urlencoded",
-        }
-        galileo_login_response = litellm.module_level_client.post(
+    def _is_configured(self) -> bool:
+        if not self.project_id or not self.base_url:
+            return False
+        if self.use_v2_api:
+            return bool(self.api_key)
+        return bool(self.username and self.password)
+
+    async def async_set_galileo_headers(self) -> None:
+        galileo_login_response = await self.async_httpx_handler.post(
             url=f"{self.base_url}/login",
-            headers=headers,
+            headers={
+                "accept": "application/json",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
             data={
-                "username": os.getenv("GALILEO_USERNAME"),
-                "password": os.getenv("GALILEO_PASSWORD"),
+                "username": self.username,
+                "password": self.password,
             },
         )
-
+        galileo_login_response.raise_for_status()
         access_token = galileo_login_response.json()["access_token"]
-
         self.headers = {
             "accept": "application/json",
             "Content-Type": "application/json",
             "Authorization": f"Bearer {access_token}",
         }
 
-    def get_output_str_from_response(self, response_obj, kwargs):
-        output = None
-        if response_obj is not None and (
-            kwargs.get("call_type", None) == "embedding"
-            or isinstance(response_obj, litellm.EmbeddingResponse)
-        ):
-            output = None
-        elif response_obj is not None and isinstance(
-            response_obj, litellm.ModelResponse
-        ):
-            output = response_obj["choices"][0]["message"].json()
-        elif response_obj is not None and isinstance(
-            response_obj, litellm.TextCompletionResponse
-        ):
-            output = response_obj.choices[0].text
-        elif response_obj is not None and isinstance(
-            response_obj, litellm.ImageResponse
-        ):
-            output = response_obj["data"]
+    async def _ensure_headers(self) -> bool:
+        if self.headers is not None:
+            return True
+
+        if self.use_v2_api:
+            if not self.api_key:
+                return False
+            self.headers = {
+                "accept": "application/json",
+                "Content-Type": "application/json",
+                "Galileo-API-Key": self.api_key,
+            }
+            return True
+
+        if not (self.username and self.password and self.base_url):
+            return False
 
-        return output
+        try:
+            await self.async_set_galileo_headers()
+            return True
+        except Exception as e:
+            verbose_logger.debug("Galileo Logger: failed to authenticate: %s", e)
+            return False
+
+    @staticmethod
+    def _galileo_input_messages(
+        messages: Optional[List[Any]], input_text: str
+    ) -> List[Dict[str, str]]:
+        if not messages:
+            return [{"role": "user", "content": input_text}]
+
+        galileo_messages: List[Dict[str, str]] = []
+        for message in messages:
+            if not isinstance(message, dict):
+                continue
+            role = message.get("role")
+            if not role:
+                continue
+            galileo_messages.append(
+                {
+                    "role": str(role),
+                    "content": convert_content_list_to_str(message=message),
+                }
+            )
+
+        if galileo_messages:
+            return galileo_messages
+        return [{"role": "user", "content": input_text}]
+
+    @staticmethod
+    def _record_to_v2_span(record: Dict[str, Any]) -> Dict[str, Any]:
+        created_at = record.get("created_at", "")
+        if created_at and "Z" not in created_at and "+" not in created_at:
+            created_at = f"{created_at}Z"
+
+        span: Dict[str, Any] = {
+            "type": "llm",
+            "name": record.get("node_type", "litellm"),
+            "created_at": created_at,
+            "input": GalileoObserve._galileo_input_messages(
+                record.get("messages"), record.get("input_text", "")
+            ),
+            "output": {
+                "role": "assistant",
+                "content": record.get("output_text", ""),
+            },
+            "status_code": record.get("status_code", 200),
+            "model": record.get("model"),
+            "metrics": {
+                "duration_ns": int(record.get("latency_ms", 0)) * 1_000_000,
+                "num_input_tokens": record.get("num_input_tokens"),
+                "num_output_tokens": record.get("num_output_tokens"),
+            },
+        }
+        if record.get("tags"):
+            span["tags"] = record["tags"]
+        return span
+
+    def _get_ingest_request(self) -> Optional[Tuple[str, Dict[str, Any]]]:
+        if not self.base_url or not self.project_id:
+            return None
+
+        if self.use_v2_api:
+            payload: Dict[str, Any] = {
+                "spans": [
+                    self._record_to_v2_span(record) for record in self.in_memory_records
+                ],
+                "reliable": False,
+            }
+            if self.log_stream_id:
+                payload["log_stream_id"] = self.log_stream_id
+            return (
+                f"{self.base_url}/v2/projects/{self.project_id}/spans",
+                payload,
+            )
+
+        return (
+            f"{self.base_url}/projects/{self.project_id}/observe/ingest",
+            {"records": self.in_memory_records},
+        )
+
+    def get_output_str_from_response(
+        self, response_obj: Any, kwargs: Dict[str, Any]
+    ) -> Optional[str]:
+        if response_obj is None:
+            return None
+        if kwargs.get("call_type", None) == "embedding" or isinstance(
+            response_obj, litellm.EmbeddingResponse
+        ):
+            return None
+        if isinstance(response_obj, litellm.TextCompletionResponse):
+            return response_obj.choices[0].text
+        if isinstance(response_obj, litellm.ImageResponse):
+            return json.dumps(response_obj["data"])
+        if isinstance(response_obj, (litellm.ModelResponse, dict)):
+            return get_content_from_model_response(response_obj)
+        return None
 
     async def async_log_success_event(
         self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any
     ):
         verbose_logger.debug("On Async Success")
 
+        if not self._is_configured():
+            verbose_logger.debug(
+                "Galileo Logger: skipping flush — set GALILEO_PROJECT_ID and "
+                "either GALILEO_API_KEY (hosted) or GALILEO_USERNAME/GALILEO_PASSWORD "
+                "(enterprise Observe)."
+            )
+            return
+
         _latency_ms = int((end_time - start_time).total_seconds() * 1000)
         _call_type = kwargs.get("call_type", "litellm")
         input_text = litellm.utils.get_formatted_prompt(
@@ -125,24 +250,48 @@ async def async_log_success_event(
                 ),  # timestamp str constructed in "%Y-%m-%dT%H:%M:%S" format
             )
 
-            # dump to dict
             request_dict = request_record.model_dump()
+            messages = kwargs.get("messages")
+            if messages:
+                request_dict["messages"] = messages
             self.in_memory_records.append(request_dict)
 
             if len(self.in_memory_records) >= self.batch_size:
                 await self.flush_in_memory_records()
 
     async def flush_in_memory_records(self):
-        verbose_logger.debug("flushing in memory records")
-        response = await self.async_httpx_handler.post(
-            url=f"{self.base_url}/projects/{self.project_id}/observe/ingest",
-            headers=self.headers,
-            json={"records": self.in_memory_records},
-        )
+        if not self.in_memory_records:
+            return
+
+        ingest_request = self._get_ingest_request()
+        if ingest_request is None:
+            verbose_logger.debug(
+                "Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID"
+            )
+            return
+
+        if not await self._ensure_headers():
+            verbose_logger.debug("Galileo Logger: could not set request headers")
+            return
+
+        url, payload = ingest_request
+        verbose_logger.debug("flushing in memory records to %s", url)
+
+        try:
+            response = await self.async_httpx_handler.post(
+                url=url,
+                headers=self.headers,
+                json=payload,
+            )
+        except Exception as e:
+            verbose_logger.debug(
+                "Galileo Logger: failed to flush in memory records: %s", e
+            )
+            return
 
         if response.status_code == 200:
             verbose_logger.debug(
-                "Galileo Logger:successfully flushed in memory records"
+                "Galileo Logger: successfully flushed in memory records"
             )
             self.in_memory_records = []
         else: