DSPy enhancements (#362)

* fix dspy issue

* DSPy enhancements v1

* Add support for LiteLLM

* update readme

* DSPy enhancements

* bump version

Author: karthikscale3

README.md

@@ -238,6 +238,14 @@ By default, prompt and completion data are captured. If you would like to opt ou
 
 `TRACE_PROMPT_COMPLETION_DATA=false`
 
+### Enable/Disable checkpoint tracing for DSPy
+
+By default, checkpoints are traced for DSPy pipelines. If you would like to disable it, set the following env var,
+
+`TRACE_DSPY_CHECKPOINT=false`
+
+Note: Checkpoint tracing will increase the latency of executions as the state is serialized. Please disable it in production.
+
 ## Supported integrations
 
 Langtrace automatically captures traces from the following vendors:
@@ -253,8 +261,9 @@ Langtrace automatically captures traces from the following vendors:
 | Gemini        | LLM             | :x:                | :white_check_mark:              |
 | Mistral       | LLM             | :x:                | :white_check_mark:              |
 | Langchain     | Framework       | :x:                | :white_check_mark:              |
-| LlamaIndex    | Framework       | :white_check_mark: | :white_check_mark:              |
 | Langgraph     | Framework       | :x:                | :white_check_mark:              |
+| LlamaIndex    | Framework       | :white_check_mark: | :white_check_mark:              |
+| LiteLLM       | Framework       | :x:                | :white_check_mark:              |
 | DSPy          | Framework       | :x:                | :white_check_mark:              |
 | CrewAI        | Framework       | :x:                | :white_check_mark:              |
 | Ollama        | Framework       | :x:                | :white_check_mark:              |

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
   'fsspec>=2024.6.0',
   "transformers>=4.11.3",
   "sentry-sdk>=2.14.0",
+  "ujson>=5.10.0",
 ]
 
 requires-python = ">=3.9"
@@ -47,6 +48,7 @@ dev = [
   "langchain-community",
   "langchain-openai",
   "langchain-openai",
+  "litellm",
   "chromadb",
   "cohere",
   "qdrant_client",

src/examples/dspy_example/optimizers/bootstrap_fewshot.py

@@ -0,0 +1,89 @@
+import dspy
+from dotenv import find_dotenv, load_dotenv
+from dspy.datasets import HotPotQA
+from dspy.teleprompt import BootstrapFewShot
+
+from langtrace_python_sdk import inject_additional_attributes, langtrace
+
+_ = load_dotenv(find_dotenv())
+
+langtrace.init()
+
+turbo = dspy.LM('openai/gpt-4o-mini')
+colbertv2_wiki17_abstracts = dspy.ColBERTv2(url='http://20.102.90.50:2017/wiki17_abstracts')
+
+dspy.settings.configure(lm=turbo, rm=colbertv2_wiki17_abstracts)
+
+
+# Load the dataset.
+dataset = HotPotQA(train_seed=1, train_size=20, eval_seed=2023, dev_size=50, test_size=0)
+
+# Tell DSPy that the 'question' field is the input. Any other fields are labels and/or metadata.
+trainset = [x.with_inputs('question') for x in dataset.train]
+devset = [x.with_inputs('question') for x in dataset.dev]
+
+
+class GenerateAnswer(dspy.Signature):
+    """Answer questions with short factoid answers."""
+
+    context = dspy.InputField(desc="may contain relevant facts")
+    question = dspy.InputField()
+    answer = dspy.OutputField(desc="often between 1 and 5 words")
+
+
+class RAG(dspy.Module):
+    def __init__(self, num_passages=3):
+        super().__init__()
+
+        self.retrieve = dspy.Retrieve(k=num_passages)
+        self.generate_answer = dspy.ChainOfThought(GenerateAnswer)
+
+    def forward(self, question):
+        context = self.retrieve(question).passages
+        prediction = self.generate_answer(context=context, question=question)
+        return dspy.Prediction(context=context, answer=prediction.answer)
+
+
+# Validation logic: check that the predicted answer is correct.
+# Also check that the retrieved context does actually contain that answer.
+def validate_context_and_answer(example, prediction, trace=None):
+    answer_em = dspy.evaluate.answer_exact_match(example, prediction)
+    answer_pm = dspy.evaluate.answer_passage_match(example, prediction)
+    return answer_em and answer_pm
+
+
+# Set up a basic optimizer, which will compile our RAG program.
+optimizer = BootstrapFewShot(metric=validate_context_and_answer)
+
+# Compile!
+compiled_rag = optimizer.compile(RAG(), trainset=trainset)
+
+# Ask any question you like to this simple RAG program.
+my_question = "Who was the hero of the movie peraanmai?"
+
+# Get the prediction. This contains `pred.context` and `pred.answer`.
+# pred = compiled_rag(my_question)
+pred = inject_additional_attributes(lambda: compiled_rag(my_question), {'experiment': 'experiment 6', 'description': 'trying additional stuff', 'run_id': 'run_1'})
+# compiled_rag.save('compiled_rag_v1.json')
+
+# Print the contexts and the answer.
+print(f"Question: {my_question}")
+print(f"Predicted Answer: {pred.answer}")
+print(f"Retrieved Contexts (truncated): {[c[:200] + '...' for c in pred.context]}")
+
+# print("Inspecting the history of the optimizer:")
+# turbo.inspect_history(n=1)
+
+from dspy.evaluate import Evaluate
+
+
+def validate_answer(example, pred, trace=None):
+    return True
+
+
+# Set up the evaluator, which can be used multiple times.
+evaluate = Evaluate(devset=devset, metric=validate_answer, num_threads=4, display_progress=True, display_table=0)
+
+
+# Evaluate our `optimized_cot` program.
+evaluate(compiled_rag)

src/examples/openai_example/chat_completion.py

@@ -9,19 +9,19 @@
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init(write_spans_to_console=True)
+langtrace.init()
 client = OpenAI()
 
 
 def api():
     response = client.chat.completions.create(
-        model="gpt-4",
+        model="o1-mini",
         messages=[
-            {"role": "system", "content": "Talk like a pirate"},
-            {"role": "user", "content": "Tell me a story in 3 sentences or less."},
+            # {"role": "system", "content": "Talk like a pirate"},
+            {"role": "user", "content": "How many r's are in strawberry?"},
         ],
-        stream=True,
-        # stream=False,
+        # stream=True,
+        stream=False,
     )
     return response
 
@@ -31,14 +31,17 @@ def chat_completion():
     response = api()
     # print(response)
     # Uncomment this for streaming
-    result = []
-    for chunk in response:
-        if chunk.choices[0].delta.content is not None:
-            content = [
-                choice.delta.content if choice.delta and choice.delta.content else ""
-                for choice in chunk.choices
-            ]
-            result.append(content[0] if len(content) > 0 else "")
-
-    # print("".join(result))
+    # result = []
+    # for chunk in response:
+    #     if chunk.choices[0].delta.content is not None:
+    #         content = [
+    #             choice.delta.content if choice.delta and choice.delta.content else ""
+    #             for choice in chunk.choices
+    #         ]
+    #         result.append(content[0] if len(content) > 0 else "")
+
+    # # print("".join(result))
+    print(response)
     return response
+
+chat_completion()

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -19,6 +19,7 @@
     "LANGCHAIN_COMMUNITY": "Langchain Community",
     "LANGCHAIN_CORE": "Langchain Core",
     "LANGGRAPH": "Langgraph",
+    "LITELLM": "Litellm",
     "LLAMAINDEX": "LlamaIndex",
     "OPENAI": "OpenAI",
     "PINECONE": "Pinecone",

src/langtrace_python_sdk/constants/instrumentation/litellm.py

@@ -0,0 +1,18 @@
+APIS = {
+    "CHAT_COMPLETION": {
+        "METHOD": "chat.completions.create",
+        "ENDPOINT": "/chat/completions",
+    },
+    "IMAGES_GENERATION": {
+        "METHOD": "images.generate",
+        "ENDPOINT": "/images/generations",
+    },
+    "IMAGES_EDIT": {
+        "METHOD": "images.edit",
+        "ENDPOINT": "/images/edits",
+    },
+    "EMBEDDINGS_CREATE": {
+        "METHOD": "embeddings.create",
+        "ENDPOINT": "/embeddings",
+    },
+}

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -19,6 +19,7 @@
 from .gemini import GeminiInstrumentation
 from .mistral import MistralInstrumentation
 from .embedchain import EmbedchainInstrumentation
+from .litellm import LiteLLMInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -31,6 +32,7 @@
     "LangchainCommunityInstrumentation",
     "LangchainCoreInstrumentation",
     "LanggraphInstrumentation",
+    "LiteLLMInstrumentation",
     "LlamaindexInstrumentation",
     "OpenAIInstrumentation",
     "PineconeInstrumentation",

src/langtrace_python_sdk/instrumentation/dspy/patch.py

@@ -1,6 +1,19 @@
 import json
+import os
+
+import ujson
+from colorama import Fore
 from importlib_metadata import version as v
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry import baggage
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
 from langtrace_python_sdk.utils import set_span_attribute
 from langtrace_python_sdk.utils.llm import (
     get_extra_attributes,
@@ -9,14 +22,6 @@
     set_span_attributes,
 )
 from langtrace_python_sdk.utils.silently_fail import silently_fail
-from langtrace_python_sdk.constants.instrumentation.common import (
-    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
-    SERVICE_PROVIDERS,
-)
-from opentelemetry import baggage
-from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry.trace import SpanKind
-from opentelemetry.trace.status import Status, StatusCode
 
 
 def patch_bootstrapfewshot_optimizer(operation_name, version, tracer):
@@ -115,6 +120,8 @@ def traced_method(wrapped, instance, args, kwargs):
             **get_extra_attributes(),
         }
 
+        trace_checkpoint = os.environ.get("TRACE_DSPY_CHECKPOINT", "true").lower()
+
         if instance.__class__.__name__:
             span_attributes["dspy.signature.name"] = instance.__class__.__name__
             span_attributes["dspy.signature"] = str(instance.signature)
@@ -136,6 +143,9 @@ def traced_method(wrapped, instance, args, kwargs):
                         "dspy.signature.result",
                         json.dumps(result.toDict()),
                     )
+                    if trace_checkpoint == "true":
+                        print(Fore.RED + "Note: DSPy checkpoint tracing is enabled in Langtrace. To disable it, set the env var, TRACE_DSPY_CHECKPOINT to false" + Fore.RESET)
+                        set_span_attribute(span, "dspy.checkpoint", ujson.dumps(instance.dump_state(False), indent=2))
                     span.set_status(Status(StatusCode.OK))
 
                 span.end()

src/langtrace_python_sdk/instrumentation/litellm/__init__.py

@@ -0,0 +1,5 @@
+from .instrumentation import LiteLLMInstrumentation
+
+__all__ = [
+    "LiteLLMInstrumentation",
+]

src/langtrace_python_sdk/instrumentation/litellm/instrumentation.py

@@ -0,0 +1,87 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Collection, Optional, Any
+import importlib.metadata
+import logging
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer, TracerProvider
+from wrapt import wrap_function_wrapper
+
+from langtrace_python_sdk.instrumentation.litellm.patch import (
+    async_chat_completions_create,
+    async_embeddings_create,
+    async_images_generate,
+    chat_completions_create,
+    embeddings_create,
+    images_generate,
+)
+
+logging.basicConfig(level=logging.FATAL)
+
+
+class LiteLLMInstrumentation(BaseInstrumentor): # type: ignore
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["litellm >= 1.48.0", "trace-attributes >= 4.0.5"]
+
+    def _instrument(self, **kwargs: Any) -> None:
+        tracer_provider: Optional[TracerProvider] = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version: str = importlib.metadata.version("openai")
+
+        wrap_function_wrapper(
+            "litellm",
+            "completion",
+            chat_completions_create(version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "litellm",
+            "text_completion",
+            chat_completions_create(version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "litellm.main",
+            "acompletion",
+            async_chat_completions_create(version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "litellm.main",
+            "image_generation",
+            images_generate(version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "litellm.main",
+            "aimage_generation",
+            async_images_generate(version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "litellm.main",
+            "embedding",
+            embeddings_create(version, tracer),
+        )
+
+        wrap_function_wrapper(
+            "litellm.main",
+            "aembedding",
+            async_embeddings_create(version, tracer),
+        )
+
+    def _uninstrument(self, **kwargs: Any) -> None:
+        pass