Merge branch 'main' of github.com:Scale3-Labs/langtrace-python-sdk into release

Author: karthikscale3

.vscode/settings.json

@@ -1,7 +1,11 @@
-
 {
   "[python]": {
     "editor.defaultFormatter": "ms-python.black-formatter",
   },
   "editor.formatOnSave": true,
+  "python.testing.pytestArgs": [
+    "src"
+  ],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true,
 }

mypy.ini

@@ -0,0 +1,16 @@
+[mypy]
+strict = True
+disable_error_code = import-untyped
+disallow_untyped_calls = True        # Disallow function calls without type annotations
+disallow_untyped_defs = True         # Disallow defining functions without type annotations
+disallow_any_explicit = True         # Disallow explicit use of `Any`
+disallow_any_generics = True         # Disallow generic types without specific type parameters
+disallow_incomplete_defs = True      # Disallow defining incomplete function signatures
+no_implicit_optional = True          # Disallow implicitly Optional types
+warn_unused_configs = True           # Warn about unused configurations
+warn_redundant_casts = True          # Warn about unnecessary type casts
+warn_return_any = True               # Warn if a function returns `Any`
+warn_unreachable = True              # Warn about unreachable code
+# Ignore external modules or allow specific imports
+follow_imports = skip
+ignore_missing_imports = True

src/examples/anthropic_example/completion.py

@@ -7,7 +7,7 @@
 
 _ = load_dotenv(find_dotenv())
 
-langtrace.init()
+langtrace.init(write_spans_to_console=True)
 
 
 @with_langtrace_root_span("messages_create")

src/examples/azureopenai_example/__init__.py

@@ -0,0 +1,9 @@
+from examples.azureopenai_example.completion import chat_completion
+from langtrace_python_sdk import with_langtrace_root_span, langtrace
+
+langtrace.init()
+
+class AzureOpenAIRunner:
+    @with_langtrace_root_span("AzureOpenAI")
+    def run(self):
+        chat_completion()

src/examples/azureopenai_example/completion.py

@@ -0,0 +1,22 @@
+import os
+from langchain_openai import AzureChatOpenAI
+
+from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+
+model = AzureChatOpenAI(
+    azure_endpoint=os.environ['AZURE_OPENAI_ENDPOINT'],
+    azure_deployment=os.environ['AZURE_OPENAI_DEPLOYMENT_NAME'],
+    openai_api_version=os.environ['AZURE_OPENAI_API_VERSION'],
+)
+
+@with_langtrace_root_span()
+def chat_completion():
+    messages = [
+        (
+            "system",
+            "You are a helpful assistant that translates English to French. Translate the user sentence.",
+        ),
+        ("human", "I love programming."),
+    ]
+    result = model.invoke(messages)
+    print(result)

src/examples/langchain_example/__init__.py

@@ -2,6 +2,7 @@
 from langtrace_python_sdk import with_langtrace_root_span
 
 from .groq_example import groq_basic, groq_streaming
+from .langgraph_example_tools import basic_graph_tools
 
 
 class LangChainRunner:
@@ -10,6 +11,7 @@ def run(self):
         basic_app()
         rag()
         load_and_split()
+        basic_graph_tools()
 
 
 class GroqRunner:

src/examples/langchain_example/langgraph_example_tools.py

@@ -0,0 +1,151 @@
+from typing import Annotated
+
+from langchain_anthropic import ChatAnthropic
+from langchain_core.messages import HumanMessage
+from langchain_core.pydantic_v1 import BaseModel
+from typing_extensions import TypedDict
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.tools import Tool
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_core.messages import AIMessage, ToolMessage
+
+from langtrace_python_sdk import langtrace
+
+langtrace.init()
+
+primes = {998: 7901, 999: 7907, 1000: 7919}
+
+
+class PrimeInput(BaseModel):
+    n: int = Field()
+
+
+def is_prime(n: int) -> bool:
+    if n <= 1 or (n % 2 == 0 and n > 2):
+        return False
+    for i in range(3, int(n**0.5) + 1, 2):
+        if n % i == 0:
+            return False
+    return True
+
+
+def get_prime(n: int, primes: dict = primes) -> str:
+    return str(primes.get(int(n)))
+
+
+async def aget_prime(n: int, primes: dict = primes) -> str:
+    return str(primes.get(int(n)))
+
+
+class State(TypedDict):
+    messages: Annotated[list, add_messages]
+    # This flag is new
+    ask_human: bool
+
+
+class RequestAssistance(BaseModel):
+    """Escalate the conversation to an expert. Use this if you are unable to assist directly or if the user requires support beyond your permissions.
+
+    To use this function, relay the user's 'request' so the expert can provide the right guidance.
+    """
+
+    request: str
+
+
+llm = ChatAnthropic(model="claude-3-haiku-20240307")
+# We can bind the llm to a tool definition, a pydantic model, or a json schema
+llm_with_tools = llm.bind_tools([RequestAssistance])
+tools = [
+    Tool(
+        name="GetPrime",
+        func=get_prime,
+        description="A tool that returns the `n`th prime number",
+        args_schema=PrimeInput,
+        coroutine=aget_prime,
+    ),
+]
+
+
+def chatbot(state: State):
+    response = llm_with_tools.invoke(state["messages"])
+    ask_human = False
+    if (
+        response.tool_calls
+        and response.tool_calls[0]["name"] == RequestAssistance.__name__
+    ):
+        ask_human = True
+    return {"messages": [response], "ask_human": ask_human}
+
+
+graph_builder = StateGraph(State)
+
+graph_builder.add_node("chatbot", chatbot)
+graph_builder.add_node("tools", ToolNode(tools=tools))
+
+
+def create_response(response: str, ai_message: AIMessage):
+    return ToolMessage(
+        content=response,
+        tool_call_id=ai_message.tool_calls[0]["id"],
+    )
+
+
+def human_node(state: State):
+    new_messages = []
+    if not isinstance(state["messages"][-1], ToolMessage):
+        # Typically, the user will have updated the state during the interrupt.
+        # If they choose not to, we will include a placeholder ToolMessage to
+        # let the LLM continue.
+        new_messages.append(
+            create_response("No response from human.", state["messages"][-1])
+        )
+    return {
+        # Append the new messages
+        "messages": new_messages,
+        # Unset the flag
+        "ask_human": False,
+    }
+
+
+def select_next_node(state: State):
+    if state["ask_human"]:
+        return "human"
+    # Otherwise, we can route as before
+    return tools_condition(state)
+
+
+def basic_graph_tools():
+    graph_builder.add_node("human", human_node)
+    graph_builder.add_conditional_edges(
+        "chatbot",
+        select_next_node,
+        {"human": "human", "tools": "tools", "__end__": "__end__"},
+    )
+    graph_builder.add_edge("tools", "chatbot")
+    graph_builder.add_edge("human", "chatbot")
+    graph_builder.set_entry_point("chatbot")
+    memory = MemorySaver()
+    graph = graph_builder.compile(
+        checkpointer=memory,
+        interrupt_before=["human"],
+    )
+
+    config = {"configurable": {"thread_id": "1"}}
+    events = graph.stream(
+        {
+            "messages": [
+                (
+                    "user",
+                    "I'm learning LangGraph. Could you do some research on it for me?",
+                )
+            ]
+        },
+        config,
+        stream_mode="values",
+    )
+    for event in events:
+        if "messages" in event:
+            event["messages"][-1]

src/examples/openai_example/__init__.py

@@ -20,3 +20,4 @@ def run(self):
         chat_completion_example()
         embeddings_create_example()
         function_example()
+        image_edit()

src/langtrace_python_sdk/instrumentation/anthropic/instrumentation.py

@@ -16,38 +16,39 @@
 
 import importlib.metadata
 import logging
-from typing import Collection
+from typing import Collection, Any
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import TracerProvider
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
-
+from typing import Any
 from langtrace_python_sdk.instrumentation.anthropic.patch import messages_create
 
 logging.basicConfig(level=logging.FATAL)
 
 
-class AnthropicInstrumentation(BaseInstrumentor):
+class AnthropicInstrumentation(BaseInstrumentor):  # type: ignore[misc]
     """
-    The AnthropicInstrumentation class represents the Anthropic instrumentation
+    The AnthropicInstrumentation class represents the Anthropic instrumentation.
     """
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return ["anthropic >= 0.19.1"]
 
-    def _instrument(self, **kwargs):
-        tracer_provider = kwargs.get("tracer_provider")
+    def _instrument(self, **kwargs: dict[str, Any]) -> None:
+        tracer_provider: TracerProvider = kwargs.get("tracer_provider")  # type: ignore
         tracer = get_tracer(__name__, "", tracer_provider)
         version = importlib.metadata.version("anthropic")
 
         wrap_function_wrapper(
             "anthropic.resources.messages",
             "Messages.create",
-            messages_create("anthropic.messages.create", version, tracer),
+            messages_create(version, tracer),
         )
 
-    def _instrument_module(self, module_name):
+    def _instrument_module(self, module_name: str) -> None:
         pass
 
-    def _uninstrument(self, **kwargs):
+    def _uninstrument(self, **kwargs: dict[str, Any]) -> None:
         pass