arize-phoenix 0.0.39__py3-none-any.whl → 0.0.41__py3-none-any.whl

phoenix/trace/fixtures.py

@@ -19,7 +19,7 @@ llama_index_rag_fixture = TracesFixture(
     file_name="llama_index_rag_v6.jsonl",
 )
 
-llama_index_calculator_agent = TracesFixture(
+llama_index_calculator_agent_fixture = TracesFixture(
     name="llama_index_calculator_agent",
     description="Traces from running the llama_index with calculator tools.",
     file_name="llama_index_calculator_agent_v2.jsonl",
@@ -43,6 +43,12 @@ langchain_titanic_csv_agent_evaluator_fixture = TracesFixture(
     file_name="lc_titanic.jsonl",
 )
 
+langchain_qa_with_sources_fixture = TracesFixture(
+    name="langchain_qa_with_sources",
+    description="LangChain QA with sources on financial data",
+    file_name="langchain_qa_with_sources_chain.jsonl",
+)
+
 random_fixture = TracesFixture(
     name="random",
     description="Randomly generated traces",
@@ -55,7 +61,8 @@ TRACES_FIXTURES: List[TracesFixture] = [
     langchain_rag_stuff_document_chain_fixture,
     langchain_titanic_csv_agent_evaluator_fixture,
     random_fixture,
-    llama_index_calculator_agent,
+    langchain_qa_with_sources_fixture,
+    llama_index_calculator_agent_fixture,
 ]
 
 NAME_TO_TRACES_FIXTURE = {fixture.name: fixture for fixture in TRACES_FIXTURES}

phoenix/trace/langchain/__init__.py

@@ -1,3 +1,4 @@
+from .instrumentor import LangChainInstrumentor
 from .tracer import OpenInferenceTracer
 
-__all__ = ["OpenInferenceTracer"]
+__all__ = ["OpenInferenceTracer", "LangChainInstrumentor"]

phoenix/trace/langchain/instrumentor.py

@@ -0,0 +1,37 @@
+from typing import Any, Optional
+
+from .tracer import OpenInferenceTracer
+
+
+class LangChainInstrumentor:
+    """
+    Instruments the OpenInferenceTracer for LangChain automatically by patching the
+    BaseCallbackManager in LangChain.
+    """
+
+    def __init__(self, tracer: Optional[OpenInferenceTracer] = None) -> None:
+        self._tracer = tracer if tracer is not None else OpenInferenceTracer()
+
+    def instrument(self) -> None:
+        try:
+            from langchain.callbacks.base import BaseCallbackManager
+        except ImportError:
+            # Raise a cleaner error if LangChain is not installed
+            raise ImportError(
+                "LangChain is not installed. Please install LangChain first to use the instrumentor"
+            )
+
+        source_init = BaseCallbackManager.__init__
+
+        # Keep track of the source init so we can tell if the patching occurred
+        self._source_callback_manager_init = source_init
+
+        tracer = self._tracer
+
+        # Patch the init method of the BaseCallbackManager to add the tracer
+        # to all callback managers
+        def patched_init(self: BaseCallbackManager, *args: Any, **kwargs: Any) -> None:
+            source_init(self, *args, **kwargs)
+            self.add_handler(tracer, True)
+
+        BaseCallbackManager.__init__ = patched_init  # type: ignore

phoenix/trace/llama_index/callback.py

@@ -106,14 +106,6 @@ def payload_to_semantic_attributes(
     if event_type in (CBEventType.NODE_PARSING, CBEventType.CHUNKING):
         # TODO(maybe): handle these events
         return attributes
-    if event_type == CBEventType.TEMPLATING:
-        if template := payload.get(EventPayload.TEMPLATE):
-            attributes[LLM_PROMPT_TEMPLATE] = template
-        if template_vars := payload.get(EventPayload.TEMPLATE_VARS):
-            attributes[LLM_PROMPT_TEMPLATE_VARIABLES] = template_vars
-        # TODO(maybe): other keys in the same payload
-        # EventPayload.SYSTEM_PROMPT
-        # EventPayload.QUERY_WRAPPER_PROMPT
     if EventPayload.CHUNKS in payload and EventPayload.EMBEDDINGS in payload:
         attributes[EMBEDDING_EMBEDDINGS] = [
             {EMBEDDING_TEXT: text, EMBEDDING_VECTOR: vector}
@@ -151,8 +143,6 @@ def payload_to_semantic_attributes(
             attributes.update(_get_output_messages(raw))
             if (usage := getattr(raw, "usage", None)) is not None:
                 attributes.update(_get_token_counts(usage))
-    if EventPayload.TEMPLATE in payload:
-        ...
     if event_type is CBEventType.RERANKING:
         ...  # TODO
         # if EventPayload.TOP_K in payload:
@@ -304,6 +294,20 @@ def _add_spans_to_tracer(
         parent_span_id, event_id = parent_child_id_stack.pop()
         event_data = event_id_to_event_data[event_id]
         event_type = event_data["event_type"]
+        attributes = event_data["attributes"]
+        if event_type is CBEventType.LLM:
+            while parent_child_id_stack:
+                preceding_event_parent_span_id, preceding_event_id = parent_child_id_stack[-1]
+                if preceding_event_parent_span_id != parent_span_id:
+                    break
+                preceding_event_data = event_id_to_event_data[preceding_event_id]
+                if preceding_event_data["event_type"] is not CBEventType.TEMPLATING:
+                    break
+                parent_child_id_stack.pop()
+                if payload := preceding_event_data["start_event"].payload:
+                    # Add template attributes to the LLM span to which they belong.
+                    attributes.update(_template_attributes(payload))
+
         start_event = event_data["start_event"]
         start_time = _timestamp_to_tz_aware_datetime(start_event.time)
         if event_type is CBEventType.EXCEPTION:
@@ -337,7 +341,7 @@ def _add_spans_to_tracer(
             status_code=SpanStatusCode.ERROR if span_exceptions else SpanStatusCode.OK,
             status_message="",
             parent_id=parent_span_id,
-            attributes=event_data["attributes"],
+            attributes=attributes,
             events=sorted(span_exceptions, key=lambda event: event.timestamp) or None,
             conversation=None,
         )
@@ -490,3 +494,14 @@ def _get_token_counts(usage: object) -> Iterator[Tuple[str, Any]]:
         yield LLM_TOKEN_COUNT_COMPLETION, completion_tokens
     if (total_tokens := getattr(usage, "total_tokens", None)) is not None:
         yield LLM_TOKEN_COUNT_TOTAL, total_tokens
+
+
+def _template_attributes(payload: Dict[str, Any]) -> Iterator[Tuple[str, Any]]:
+    """Yields template attributes if present"""
+    if template := payload.get(EventPayload.TEMPLATE):
+        yield LLM_PROMPT_TEMPLATE, template
+    if template_vars := payload.get(EventPayload.TEMPLATE_VARS):
+        yield LLM_PROMPT_TEMPLATE_VARIABLES, template_vars
+        # TODO(maybe): other keys in the same payload
+        # EventPayload.SYSTEM_PROMPT
+        # EventPayload.QUERY_WRAPPER_PROMPT