Merge branch 'main' into sc-27873-get-prompt-template-output

Author: shriyaedu

examples/README.md

@@ -0,0 +1,22 @@
+## Collection of Python galileo examples
+
+### Preconditions
+
+Install `uv`, we use inline dependency inside scripts.
+
+### How to use/run?
+
+First of all create `.env` file and add required env vars based on `.env.sample`.
+
+Then just run `uv`:
+
+```bash
+uv run --env-file=examples/langgraph/.env examples/langgraph/with_openai.py
+```
+
+or
+
+#### [basic_langgraph.py]
+```bash
+uv run --env-file=examples/langgraph/.env examples/langgraph/basic_langgraph.py
+```

examples/langgraph/.env.sample

@@ -0,0 +1,4 @@
+GALILEO_API_KEY=
+GALILEO_PROJECT=
+GALILEO_LOG_STREAM=
+OPENAI_API_KEY=

examples/langgraph/basic_langgraph.py

@@ -0,0 +1,48 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "galileo",
+#     "langgraph",
+#     "langsmith",
+#     "langchain",
+#     "grandalf", # for printing graph in ascii
+# ]
+# ///
+# from dotenv import load_dotenv; load_dotenv()
+
+from typing import Annotated
+
+from langchain_core.messages import AIMessage
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.message import add_messages
+from typing_extensions import TypedDict
+
+from galileo.handlers.langchain import GalileoCallback
+
+
+class State(TypedDict):
+    # Messages have the type "list". The `add_messages` function
+    # in the annotation defines how this state key should be updated
+    # (in this case, it appends messages to the list, rather than overwriting them)
+    messages: Annotated[list, add_messages]
+
+
+def node(state: State):
+    messages = state["messages"]
+    new_message = AIMessage("Hello!")
+
+    return {"messages": messages + [new_message], "extra_field": 10}
+
+
+def node2(state: State):
+    return {"messages": state["messages"]}
+
+
+graph_builder = StateGraph(State)
+graph_builder.add_node("node_name", node)
+graph_builder.add_edge(START, "node_name")
+graph_builder.add_edge("node_name", END)
+graph = graph_builder.compile()
+
+graph.get_graph().print_ascii()
+graph.invoke({"messages": [{"role": "user", "content": "hi!"}]}, config={"callbacks": [GalileoCallback()]})

examples/langgraph/with_openai.py

@@ -0,0 +1,45 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "galileo",
+#     "langgraph",
+#     "langsmith",
+#     "langchain[openai]",
+#     "grandalf", # for printing graph in ascii
+# ]
+# ///
+from typing import Annotated
+
+from langchain_openai import ChatOpenAI
+from langgraph.graph import START, StateGraph
+from langgraph.graph.message import add_messages
+from typing_extensions import TypedDict
+
+from galileo.handlers.langchain import GalileoCallback
+
+
+class State(TypedDict):
+    # Messages have the type "list". The `add_messages` function
+    # in the annotation defines how this state key should be updated
+    # (in this case, it appends messages to the list, rather than overwriting them)
+    messages: Annotated[list, add_messages]
+
+
+llm = ChatOpenAI(model="gpt-4")
+
+
+def chatbot(state: State):
+    return {"messages": [llm.invoke(state["messages"])]}
+
+
+graph_builder = StateGraph(State)
+# The first argument is the unique node name
+# The second argument is the function or object that will be called whenever
+# the node is used.
+graph_builder.add_node("chatbot", chatbot)
+graph_builder.add_edge(START, "chatbot")
+graph = graph_builder.compile()
+
+graph.get_graph().print_ascii()
+
+graph.invoke({"messages": [{"role": "user", "content": "hi!"}]}, {"callbacks": [GalileoCallback()]})

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "galileo"
-version = "0.10.0"
+version = "1.0.0"
 description = "Client library for the Galileo platform."
 authors = [{ name = "Galileo Technologies Inc.", email = "team@galileo.ai" }]
 readme = "README.md"

src/galileo/__init__.py

@@ -13,4 +13,4 @@
 from galileo_core.schemas.logging.step import StepType
 from galileo_core.schemas.logging.trace import Trace
 
-__version__ = "0.10.0"
+__version__ = "1.0.0"

src/galileo/handlers/langchain/async_handler.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import time
+from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -109,6 +110,7 @@ async def _log_node_tree(self, node: Node) -> None:
         name = node.span_params.get("name")
         metadata = node.span_params.get("metadata")
         tags = node.span_params.get("tags")
+        created_at = node.span_params.get("created_at")
 
         # Convert metadata to a dict[str, str]
         if metadata is not None:
@@ -123,6 +125,7 @@ async def _log_node_tree(self, node: Node) -> None:
                 duration_ns=node.span_params.get("duration_ns"),
                 metadata=metadata,
                 tags=tags,
+                created_at=created_at,
             )
             is_workflow_span = True
         elif node.node_type in ("llm", "chat"):
@@ -140,6 +143,7 @@ async def _log_node_tree(self, node: Node) -> None:
                 num_output_tokens=node.span_params.get("num_output_tokens"),
                 total_tokens=node.span_params.get("total_tokens"),
                 time_to_first_token_ns=node.span_params.get("time_to_first_token_ns"),
+                created_at=created_at,
             )
         elif node.node_type == "retriever":
             self._galileo_logger.add_retriever_span(
@@ -149,6 +153,7 @@ async def _log_node_tree(self, node: Node) -> None:
                 duration_ns=node.span_params.get("duration_ns"),
                 metadata=metadata,
                 tags=tags,
+                created_at=created_at,
             )
         elif node.node_type == "tool":
             self._galileo_logger.add_tool_span(
@@ -158,6 +163,7 @@ async def _log_node_tree(self, node: Node) -> None:
                 duration_ns=node.span_params.get("duration_ns"),
                 metadata=metadata,
                 tags=tags,
+                created_at=created_at,
             )
         else:
             _logger.warning(f"Unknown node type: {node.node_type}")
@@ -208,9 +214,13 @@ async def _start_node(
         # Create new node
         node = Node(node_type=node_type, span_params=kwargs, run_id=run_id, parent_run_id=parent_run_id)
 
+        # start_time is used to calculate duration_ns
         if "start_time" not in node.span_params:
             node.span_params["start_time"] = time.perf_counter_ns()
 
+        if "created_at" not in node.span_params:
+            node.span_params["created_at"] = datetime.now(tz=timezone.utc)
+
         self._nodes[node_id] = node
 
         # Set as root node if needed

src/galileo/handlers/langchain/handler.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import time
+from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -111,6 +112,7 @@ def _log_node_tree(self, node: Node) -> None:
         name = node.span_params.get("name")
         metadata = node.span_params.get("metadata")
         tags = node.span_params.get("tags")
+        created_at = node.span_params.get("created_at")
 
         # Convert metadata to a dict[str, str]
         if metadata is not None:
@@ -125,6 +127,7 @@ def _log_node_tree(self, node: Node) -> None:
                 duration_ns=node.span_params.get("duration_ns"),
                 metadata=metadata,
                 tags=tags,
+                created_at=created_at,
             )
             is_workflow_span = True
         elif node.node_type in ("llm", "chat"):
@@ -142,6 +145,7 @@ def _log_node_tree(self, node: Node) -> None:
                 num_output_tokens=node.span_params.get("num_output_tokens"),
                 total_tokens=node.span_params.get("total_tokens"),
                 time_to_first_token_ns=node.span_params.get("time_to_first_token_ns"),
+                created_at=created_at,
             )
         elif node.node_type == "retriever":
             self._galileo_logger.add_retriever_span(
@@ -151,6 +155,7 @@ def _log_node_tree(self, node: Node) -> None:
                 duration_ns=node.span_params.get("duration_ns"),
                 metadata=metadata,
                 tags=tags,
+                created_at=created_at,
             )
         elif node.node_type == "tool":
             self._galileo_logger.add_tool_span(
@@ -160,6 +165,7 @@ def _log_node_tree(self, node: Node) -> None:
                 duration_ns=node.span_params.get("duration_ns"),
                 metadata=metadata,
                 tags=tags,
+                created_at=created_at,
             )
         else:
             _logger.warning(f"Unknown node type: {node.node_type}")
@@ -210,9 +216,13 @@ def _start_node(
         # Create new node
         node = Node(node_type=node_type, span_params=kwargs, run_id=run_id, parent_run_id=parent_run_id)
 
+        # start_time is used to calculate duration_ns
         if "start_time" not in node.span_params:
             node.span_params["start_time"] = time.perf_counter_ns()
 
+        if "created_at" not in node.span_params:
+            node.span_params["created_at"] = datetime.now(tz=timezone.utc)
+
         self._nodes[node_id] = node
 
         # Set as root node if needed

tests/test_langchain.py

@@ -1,3 +1,4 @@
+import time
 import uuid
 from unittest.mock import MagicMock, Mock, patch
 
@@ -690,3 +691,59 @@ def test_callback_with_active_trace(self, galileo_logger: GalileoLogger):
         assert traces[0].spans[0].spans[0].type == "retriever"
         assert traces[0].spans[0].spans[0].input == "test query"
         assert traces[0].spans[0].spans[0].output == [GalileoDocument(content="test document", metadata={})]
+
+    def test_node_created_at(self, callback: GalileoCallback, galileo_logger: GalileoLogger):
+        parent_id = uuid.uuid4()
+        llm_run_id = uuid.uuid4()
+        retriever_run_id = uuid.uuid4()
+
+        # Create parent chain
+        callback.on_chain_start(serialized={}, inputs={"query": "test"}, run_id=parent_id)
+
+        # Start retriever
+        callback.on_retriever_start(
+            serialized={}, query="AI development", run_id=retriever_run_id, parent_run_id=parent_id
+        )
+
+        # End retriever
+        document = Document(page_content="AI is advancing rapidly", metadata={"source": "textbook"})
+        callback.on_retriever_end(documents=[document], run_id=retriever_run_id, parent_run_id=parent_id)
+
+        delay_ms = 500
+
+        time.sleep(delay_ms / 1000)
+
+        callback.on_llm_start(
+            serialized={},
+            prompts=["Tell me about AI"],
+            run_id=llm_run_id,
+            parent_run_id=parent_id,
+            invocation_params={"model_name": "gpt-4", "temperature": 0.7},
+        )
+
+        # Add a token to test token timing
+        callback.on_llm_new_token("AI", run_id=llm_run_id)
+
+        # End LLM
+        llm_response = MagicMock()
+        llm_response.generations = [[MagicMock()]]
+        llm_response.llm_output = {"token_usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}}
+
+        # Mock dict method on the generation
+        llm_response.generations[0][0].dict.return_value = {"text": "AI is a technology..."}
+
+        callback.on_llm_end(response=llm_response, run_id=llm_run_id, parent_run_id=parent_id)
+
+        # End chain
+        callback.on_chain_end(outputs='{"result": "test answer"}', run_id=parent_id)
+
+        traces = galileo_logger.traces
+        assert len(traces) == 1
+        assert len(traces[0].spans) == 1
+        assert len(traces[0].spans[0].spans) == 2
+
+        retriever_span = traces[0].spans[0].spans[0]
+        llm_span = traces[0].spans[0].spans[1]
+
+        time_diff_ms = (llm_span.created_at - retriever_span.created_at).total_seconds() * 1000
+        assert time_diff_ms >= delay_ms

tests/test_langchain_async.py

@@ -1,3 +1,4 @@
+import asyncio
 import uuid
 from unittest.mock import MagicMock, Mock, patch
 
@@ -667,3 +668,60 @@ async def test_callback_with_active_trace(self, galileo_logger: GalileoLogger):
         assert traces[0].spans[0].spans[0].type == "retriever"
         assert traces[0].spans[0].spans[0].input == "test query"
         assert traces[0].spans[0].spans[0].output == [GalileoDocument(content="test document", metadata={})]
+
+    @mark.asyncio
+    async def test_node_created_at(self, callback: GalileoAsyncCallback, galileo_logger: GalileoLogger):
+        parent_id = uuid.uuid4()
+        llm_run_id = uuid.uuid4()
+        retriever_run_id = uuid.uuid4()
+
+        # Create parent chain
+        await callback.on_chain_start(serialized={}, inputs={"query": "test"}, run_id=parent_id)
+
+        # Start retriever
+        await callback.on_retriever_start(
+            serialized={}, query="AI development", run_id=retriever_run_id, parent_run_id=parent_id
+        )
+
+        # End retriever
+        document = Document(page_content="AI is advancing rapidly", metadata={"source": "textbook"})
+        await callback.on_retriever_end(documents=[document], run_id=retriever_run_id, parent_run_id=parent_id)
+
+        delay_ms = 500
+
+        await asyncio.sleep(delay_ms / 1000)
+
+        await callback.on_llm_start(
+            serialized={},
+            prompts=["Tell me about AI"],
+            run_id=llm_run_id,
+            parent_run_id=parent_id,
+            invocation_params={"model_name": "gpt-4", "temperature": 0.7},
+        )
+
+        # Add a token to test token timing
+        await callback.on_llm_new_token("AI", run_id=llm_run_id)
+
+        # End LLM
+        llm_response = MagicMock()
+        llm_response.generations = [[MagicMock()]]
+        llm_response.llm_output = {"token_usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}}
+
+        # Mock dict method on the generation
+        llm_response.generations[0][0].dict.return_value = {"text": "AI is a technology..."}
+
+        await callback.on_llm_end(response=llm_response, run_id=llm_run_id, parent_run_id=parent_id)
+
+        # End chain
+        await callback.on_chain_end(outputs='{"result": "test answer"}', run_id=parent_id)
+
+        traces = galileo_logger.traces
+        assert len(traces) == 1
+        assert len(traces[0].spans) == 1
+        assert len(traces[0].spans[0].spans) == 2
+
+        retriever_span = traces[0].spans[0].spans[0]
+        llm_span = traces[0].spans[0].spans[1]
+
+        time_diff_ms = (llm_span.created_at - retriever_span.created_at).total_seconds() * 1000
+        assert time_diff_ms >= delay_ms