Add langchain including vectordb instrumentation (#21)

Add LangChain Instrumentation
In this Pull Request, we have added a new callback handler, NewRelicCallbackHandler, that extends the BaseCallbackHandler from the LangChain package. This new handler integrates the New Relic monitoring service with the LangChain application to provide detailed insights and observations during the lifecycle of various processes like Language Model (LLM) start/end, Chat Model start, token generation, tool start/end, chain start/end, agent action, and agent finish.

The NewRelicCallbackHandler includes methods to handle different events in the lifecycle, each associated with creating, updating, and finishing spans which help in tracing the operations in the New Relic dashboard. For example, on_llm_start method is triggered when the LLM starts running, and on_llm_end method when the LLM ends running, each creating and finishing a span respectively.

.gitignore

@@ -1,3 +1,7 @@
+.DS_store
+.vscode/
+dependencies/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -20,6 +24,7 @@ parts/
 sdist/
 var/
 wheels/
+pip-wheel-metadata/
 share/python-wheels/
 *.egg-info/
 .installed.cfg
@@ -83,9 +88,7 @@ profile_default/
 ipython_config.py
 
 # pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
+.python-version
 
 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
@@ -148,3 +151,5 @@ poetry.lock
 
 #git
 *.patch
+# poetry
+poetry.lock

download-dependencies.sh

@@ -0,0 +1,19 @@
+#!/bin/sh
+
+set -e
+
+echo "(Re)-creating directory"
+rm -rf ./dependencies
+mkdir ./dependencies
+cd ./dependencies
+echo "Downloading dependencies"
+curl -sS https://d2eo22ngex1n9g.cloudfront.net/Documentation/SDK/bedrock-python-sdk.zip > sdk.zip
+echo "Unpacking dependencies"
+# (SageMaker Studio system terminals don't have `unzip` utility installed)
+if command -v unzip &> /dev/null
+then
+    unzip sdk.zip && rm sdk.zip && echo "Done"
+else
+    echo "'unzip' command not found: Trying to unzip via Python"
+    python -m zipfile -e sdk.zip . && rm sdk.zip && echo "Done"
+fi

examples/example.py

@@ -1,16 +1,18 @@
 import os
 
 import openai
+
 from nr_openai_observability import monitor
 
 monitor.initialization(
-    metadata={"environment": "development"}
+    application_name="OpenAI observability example",
+    metadata={"environment": "development"},
 )
 
 openai.api_key = os.getenv("OPENAI_API_KEY")
 openai.Completion.create(
     model="text-davinci-003",
     prompt="What is Observability?",
     max_tokens=20,
-    temperature=0 
+    temperature=0,
 )

examples/langchain_example.py

@@ -0,0 +1,61 @@
+from typing import Any, Dict, List
+
+import boto3
+from langchain.agents import AgentType, initialize_agent
+from langchain.chat_models import ChatOpenAI
+from langchain.llms.base import LLM
+from langchain.llms.bedrock import Bedrock
+from langchain.tools import Tool
+
+from nr_openai_observability.langchain_callback import NewRelicCallbackHandler
+
+
+def metadata_callback(ocj: Any) -> Dict[str, Any]:
+    return {"conversation_id": 123}
+
+
+new_relic_monitor = NewRelicCallbackHandler(
+    "LangChain observability trace", metadata_callback=metadata_callback
+)
+
+
+def math(x):
+    return 4
+
+
+tools = []
+tools.append(
+    Tool(
+        func=math,
+        name="Calculator",
+        description="useful for when you need to answer questions about math",
+    )
+)
+
+
+openai_llm = ChatOpenAI(temperature=0)
+
+
+boto_client = boto3.client("bedrock", "us-east-1")
+bedrock_llm = Bedrock(
+    model_id="anthropic.claude-v2",  # "amazon.titan-tg1-large"
+    client=boto_client,
+)
+
+
+def get_agent(llm: LLM, tools: List[Tool]):
+    return initialize_agent(
+        tools,
+        llm,
+        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+    )
+
+
+openai_agent = get_agent(openai_llm, tools)
+openai_agent.run("What is 2 + 2?", callbacks=[new_relic_monitor])
+
+bedrock_agent = get_agent(bedrock_llm, tools)
+bedrock_agent.run("What is 2 + 2?", callbacks=[new_relic_monitor])
+
+
+print("Agent run successfully!")

examples/langchain_vectordb_example.py

@@ -0,0 +1,17 @@
+from langchain.document_loaders import TextLoader
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.vectorstores import Chroma
+
+from nr_openai_observability.langchain_callback import NewRelicCallbackHandler
+
+new_relic_monitor = NewRelicCallbackHandler("LangChain observability trace")
+
+raw_documents = TextLoader("./examples/state_of_the_union.txt").load()
+text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+documents = text_splitter.split_documents(raw_documents)
+db = Chroma.from_documents(documents, OpenAIEmbeddings())
+
+query = "What did the president say about Ketanji Brown Jackson"
+docs = db.similarity_search(query)
+print(docs[0].page_content)