refactor: address missing issues

Author: clementsirieix

examples/langchain_toolcall.py

@@ -19,6 +19,7 @@
 tools = [search]
 
 lai_client = LiteralClient()
+lai_client.initialize()
 lai_prompt = lai_client.api.get_or_create_prompt(
     name="LC Agent",
     settings={
@@ -37,13 +38,13 @@
         {"role": "assistant", "content": "{{agent_scratchpad}}"},
     ],
 )
-prompt = lai_prompt.to_langchain_chat_prompt_template()
+prompt = lai_prompt.to_langchain_chat_prompt_template(
+    additional_messages=[("placeholder", "{agent_scratchpad}")],
+)
 
 agent: BaseSingleActionAgent = create_tool_calling_agent(model, tools, prompt)  # type: ignore
 agent_executor = AgentExecutor(agent=agent, tools=tools)
 
-cb = lai_client.langchain_callback()
-
 # Replace with ainvoke for asynchronous execution.
 agent_executor.invoke(
     {
@@ -56,5 +57,5 @@
         ],
         "input": "whats the weather in sf?",
     },
-    config=RunnableConfig(callbacks=[cb], run_name="Weather SF"),
+    config=RunnableConfig(run_name="Weather SF"),
 )

examples/langchain_variable.py

@@ -1,12 +1,13 @@
 from langchain.chat_models import init_chat_model
 from literalai import LiteralClient
-from langchain.schema.runnable.config import RunnableConfig
+
 
 from dotenv import load_dotenv
 
 load_dotenv()
 
 lai = LiteralClient()
+lai.initialize()
 
 prompt = lai.api.get_or_create_prompt(
     name="user intent",
@@ -29,13 +30,14 @@
 input_messages = messages.format_messages(
     user_message="The screen is cracked, there are scratches on the surface, and a component is missing."
 )
-cb = lai.langchain_callback()
 
 # Returns a langchain_openai.ChatOpenAI instance.
 gpt_4o = init_chat_model(  # type: ignore
     model_provider=prompt.provider,
     **prompt.settings,
 )
-print(gpt_4o.invoke(input_messages, config=RunnableConfig(callbacks=[cb])))
+
+lai.set_properties(prompt=prompt)
+print(gpt_4o.invoke(input_messages))
 
 lai.flush_and_stop()

examples/multimodal.py

@@ -0,0 +1,84 @@
+import base64
+import requests  # type: ignore
+import time
+
+from literalai import LiteralClient
+from openai import OpenAI
+
+from dotenv import load_dotenv
+
+from literalai.observability.step import ScoreDict
+
+load_dotenv()
+
+openai_client = OpenAI()
+
+literalai_client = LiteralClient()
+literalai_client.initialize()
+
+
+def encode_image(url):
+    return base64.b64encode(requests.get(url).content)
+
+
+@literalai_client.step(type="run")
+def generate_answer(user_query, image_url):
+    literalai_client.set_properties(
+        name="foobar",
+        metadata={"foo": "bar"},
+        tags=["foo", "bar"],
+    )
+    completion = openai_client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": user_query},
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": image_url},
+                    },
+                ],
+            },
+        ],
+        max_tokens=300,
+    )
+    return completion.choices[0].message.content
+
+
+def main():
+    with literalai_client.thread(name="Meal Analyzer") as thread:
+        welcome_message = (
+            "Welcome to the meal analyzer, please upload an image of your plate!"
+        )
+        literalai_client.message(
+            content=welcome_message, type="assistant_message", name="My Assistant"
+        )
+
+        user_query = "Is this a healthy meal?"
+        user_image = "https://www.eatthis.com/wp-content/uploads/sites/4/2021/05/healthy-plate.jpg"
+        user_step = literalai_client.message(
+            content=user_query, type="user_message", name="User"
+        )
+
+        time.sleep(1)  # to make sure the user step has arrived at Literal AI
+
+        literalai_client.api.create_attachment(
+            thread_id=thread.id,
+            step_id=user_step.id,
+            name="meal_image",
+            content=encode_image(user_image),
+        )
+
+        answer = generate_answer(user_query=user_query, image_url=user_image)
+        literalai_client.message(
+            content=answer, type="assistant_message", name="My Assistant"
+        )
+
+
+main()
+# Network requests by the SDK are performed asynchronously.
+# Invoke flush_and_stop() to guarantee the completion of all requests prior to the process termination.
+# WARNING: If you run a continuous server, you should not use this method.
+literalai_client.flush_and_stop()

examples/streaming.py

@@ -12,7 +12,7 @@
 
 
 sdk = LiteralClient(batch_size=2)
-sdk.instrument_openai()
+sdk.initialize()
 
 
 @sdk.thread

literalai/client.py

@@ -29,6 +29,7 @@
     step_decorator,
 )
 from literalai.observability.thread import ThreadContextManager, thread_decorator
+from literalai.prompt_engineering.prompt import Prompt
 from literalai.requirements import check_all_requirements
 
 
@@ -373,6 +374,7 @@ def set_properties(
         name: Optional[str] = None,
         tags: Optional[List[str]] = None,
         metadata: Optional[Dict[str, Any]] = None,
+        prompt: Optional[Prompt] = None,
     ):
         thread = active_thread_var.get()
         root_run = active_root_run_var.get()
@@ -386,6 +388,7 @@ def set_properties(
                 "literal.name": str(name) if name else "None",
                 "literal.tags": json.dumps(tags) if tags else "None",
                 "literal.metadata": json.dumps(metadata) if metadata else "None",
+                "literal.prompt": json.dumps(prompt.to_dict()) if prompt else "None",
             }
         )
 

literalai/event_processor.py

@@ -98,7 +98,6 @@ def _process_batch(self, batch: List):
             self.processing_counter -= len(batch)
 
     def flush_and_stop(self):
-        time.sleep(4)
         self.stop_event.set()
         if not self.disabled:
             self.processing_thread.join()

literalai/exporter.py

@@ -1,5 +1,6 @@
-from datetime import datetime, timezone
+from datetime import date, datetime, timezone
 import json
+from annotated_types import Timezone
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from typing import Dict, List, Optional, Sequence, cast
@@ -10,6 +11,7 @@
 from literalai.helper import utc_now
 from literalai.observability.generation import GenerationType
 from literalai.observability.step import Step, StepDict
+from literalai.prompt_engineering.prompt import PromptDict
 
 
 class LoggingSpanExporter(SpanExporter):
@@ -56,14 +58,8 @@ def force_flush(self, timeout_millis: float = 30000) -> bool:
         """Force flush the exporter."""
         return True
 
-    #     # TODO: Add generation promptid
-    #     # TODO: Add generation variables
-    #     # TODO: Check missing variables
-    #     # TODO: ttFirstToken
-    #     # TODO: duration
-    #     # TODO: tokenThroughputInSeconds
-    #     # TODO: Add tools
-    #     # TODO: error check with gemini error
+    # TODO: error check with gemini error
+    # TODO: ttFirstToken
     def _create_step_from_span(self, span: ReadableSpan) -> Step:
         """Convert a span to a Step object"""
         attributes = span.attributes or {}
@@ -78,6 +74,11 @@ def _create_step_from_span(self, span: ReadableSpan) -> Step:
             if span.end_time
             else utc_now()
         )
+        duration, token_throughput = self._calculate_duration_and_throughput(
+            span.start_time,
+            span.end_time,
+            int(str(attributes.get("llm.usage.total_tokens", 0))),
+        )
 
         generation_type = attributes.get("llm.request.type")
         is_chat = generation_type == "chat"
@@ -103,19 +104,35 @@ def _create_step_from_span(self, span: ReadableSpan) -> Step:
             k: str(v) for k, v in span_props.items() if v is not None and v != "None"
         }
 
+        serialized_prompt = attributes.get(
+            "traceloop.association.properties.literal.prompt"
+        )
+        prompt = cast(
+            Optional[PromptDict],
+            (
+                self._extract_json(str(serialized_prompt))
+                if serialized_prompt and serialized_prompt != "None"
+                else None
+            ),
+        )
+
         generation_content = {
+            "duration": duration,
             "messages": (
-                self.extract_messages(cast(Dict, attributes)) if is_chat else None
+                self._extract_messages(cast(Dict, attributes)) if is_chat else None
             ),
             "message_completion": (
-                self.extract_messages(cast(Dict, attributes), "gen_ai.completion.")[0]
+                self._extract_messages(cast(Dict, attributes), "gen_ai.completion.")[0]
                 if is_chat
                 else None
             ),
             "prompt": attributes.get("gen_ai.prompt.0.user"),
+            "promptId": prompt.get("id") if prompt else None,
             "completion": attributes.get("gen_ai.completion.0.content"),
             "model": attributes.get("gen_ai.request.model"),
             "provider": attributes.get("gen_ai.system"),
+            "tokenThroughputInSeconds": token_throughput,
+            "variables": prompt.get("variables") if prompt else None,
         }
         generation_settings = {
             "max_tokens": attributes.get("gen_ai.request.max_tokens"),
@@ -133,13 +150,13 @@ def _create_step_from_span(self, span: ReadableSpan) -> Step:
             "id": str(span.context.span_id) if span.context else None,
             "name": span_props.get("name", span.name),
             "type": "llm",
-            "metadata": self.extract_json(span_props.get("metadata", "{}")),
+            "metadata": self._extract_json(span_props.get("metadata", "{}")),
             "startTime": start_time,
             "endTime": end_time,
             "threadId": span_props.get("thread_id"),
             "parentId": span_props.get("parent_id"),
             "rootRunId": span_props.get("root_run_id"),
-            "tags": self.extract_json(span_props.get("tags", "[]")),
+            "tags": self._extract_json(span_props.get("tags", "[]")),
             "input": {
                 "content": (
                     generation_content["messages"]
@@ -176,7 +193,7 @@ def _create_step_from_span(self, span: ReadableSpan) -> Step:
 
         return step
 
-    def extract_messages(
+    def _extract_messages(
         self, data: Dict, prefix: str = "gen_ai.prompt."
     ) -> List[Dict]:
         messages = []
@@ -188,22 +205,42 @@ def extract_messages(
 
             if role_key not in data or content_key not in data:
                 break
+            if data[role_key] == "placeholder":
+                break
 
             messages.append(
                 {
                     "role": data[role_key],
-                    "content": self.extract_json(data[content_key]),
+                    "content": self._extract_json(data[content_key]),
                 }
             )
 
             index += 1
 
         return messages
 
-    def extract_json(self, data: str) -> Dict | List | str:
+    def _extract_json(self, data: str) -> Dict | List | str:
         try:
             content = json.loads(data)
         except Exception:
             content = data
 
         return content
+
+    def _calculate_duration_and_throughput(
+        self,
+        start_time_ns: Optional[int],
+        end_time_ns: Optional[int],
+        total_tokens: Optional[int],
+    ) -> tuple[float, Optional[float]]:
+        """Calculate duration in seconds and token throughput per second."""
+        duration_ns = (
+            end_time_ns - start_time_ns if start_time_ns and end_time_ns else 0
+        )
+        duration_seconds = duration_ns / 1e9
+
+        token_throughput = None
+        if total_tokens is not None and duration_seconds > 0:
+            token_throughput = total_tokens / duration_seconds
+
+        return duration_seconds, token_throughput