Skip to content

Commit ced2e9f

Browse files
gustavocidornelaswhoseoyster
gustavocidornelas
authored and
whoseoyster
committed
feat: support Ollama models via LangChain callback handler
1 parent 3920d2a commit ced2e9f

File tree

2 files changed

+177
-8
lines changed

2 files changed

+177
-8
lines changed

‎examples/tracing/ollama/ollama_tracing.ipynb‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "2722b419",
6+
"metadata": {},
7+
"source": [
8+
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openlayer-ai/openlayer-python/blob/main/examples/tracing/ollama/ollama_tracing.ipynb)\n",
9+
"\n",
10+
"\n",
11+
"# <a id=\"top\">Ollama tracing</a>\n",
12+
"\n",
13+
"This notebook illustrates how use Openlayer's callback handler to trace Ollama calls. \n",
14+
"\n",
15+
"Before running this notebook, make sure you first follow [these instructions](https://github.com/ollama/ollama) to set up and run a local Ollama instance."
16+
]
17+
},
18+
{
19+
"cell_type": "code",
20+
"execution_count": null,
21+
"id": "020c8f6a",
22+
"metadata": {},
23+
"outputs": [],
24+
"source": [
25+
"!pip install openlayer langchain-ollama"
26+
]
27+
},
28+
{
29+
"cell_type": "markdown",
30+
"id": "75c2a473",
31+
"metadata": {},
32+
"source": [
33+
"## 1. Set the environment variables"
34+
]
35+
},
36+
{
37+
"cell_type": "code",
38+
"execution_count": null,
39+
"id": "f3f4fa13",
40+
"metadata": {},
41+
"outputs": [],
42+
"source": [
43+
"import os\n",
44+
"\n",
45+
"# Openlayer env variables\n",
46+
"os.environ[\"OPENLAYER_API_KEY\"] = \"YOUR_OPENLAYER_API_KEY_HERE\"\n",
47+
"os.environ[\"OPENLAYER_INFERENCE_PIPELINE_ID\"] = \"YOUR_OPENLAYER_INFERENCE_PIPELINE_ID_HERE\""
48+
]
49+
},
50+
{
51+
"cell_type": "markdown",
52+
"id": "9758533f",
53+
"metadata": {},
54+
"source": [
55+
"## 2. Instantiate the `OpenlayerHandler`"
56+
]
57+
},
58+
{
59+
"cell_type": "code",
60+
"execution_count": null,
61+
"id": "e60584fa",
62+
"metadata": {},
63+
"outputs": [],
64+
"source": [
65+
"from openlayer.lib.integrations import langchain_callback\n",
66+
"\n",
67+
"openlayer_handler = langchain_callback.OpenlayerHandler()"
68+
]
69+
},
70+
{
71+
"cell_type": "markdown",
72+
"id": "76a350b4",
73+
"metadata": {},
74+
"source": [
75+
"## 3. Use an Ollama model with LangChain\n",
76+
"\n",
77+
"Now, you can pass the `openlayer_handler` as a callback to LLM's or chain invokations."
78+
]
79+
},
80+
{
81+
"cell_type": "code",
82+
"execution_count": null,
83+
"id": "e00c1c79",
84+
"metadata": {},
85+
"outputs": [],
86+
"source": [
87+
"from langchain_ollama import ChatOllama"
88+
]
89+
},
90+
{
91+
"cell_type": "code",
92+
"execution_count": null,
93+
"id": "abaf6987-c257-4f0d-96e7-3739b24c7206",
94+
"metadata": {},
95+
"outputs": [],
96+
"source": [
97+
"chat = ChatOllama(\n",
98+
" model=\"llama3.1\",\n",
99+
" callbacks=[openlayer_handler]\n",
100+
")"
101+
]
102+
},
103+
{
104+
"cell_type": "code",
105+
"execution_count": null,
106+
"id": "4123669f-aa28-47b7-8d46-ee898aba99e8",
107+
"metadata": {},
108+
"outputs": [],
109+
"source": [
110+
"chat.invoke(\"What's the meaning of life?\")"
111+
]
112+
},
113+
{
114+
"cell_type": "markdown",
115+
"id": "9a702ad1-da68-4757-95a6-4661ddaef251",
116+
"metadata": {},
117+
"source": [
118+
"That's it! Now your data is being streamed to Openlayer after every invokation."
119+
]
120+
},
121+
{
122+
"cell_type": "code",
123+
"execution_count": null,
124+
"id": "a3092828-3fbd-4f12-bae7-8de7f7319ff0",
125+
"metadata": {},
126+
"outputs": [],
127+
"source": []
128+
}
129+
],
130+
"metadata": {
131+
"kernelspec": {
132+
"display_name": "Python 3 (ipykernel)",
133+
"language": "python",
134+
"name": "python3"
135+
},
136+
"language_info": {
137+
"codemirror_mode": {
138+
"name": "ipython",
139+
"version": 3
140+
},
141+
"file_extension": ".py",
142+
"mimetype": "text/x-python",
143+
"name": "python",
144+
"nbconvert_exporter": "python",
145+
"pygments_lexer": "ipython3",
146+
"version": "3.9.19"
147+
}
148+
},
149+
"nbformat": 4,
150+
"nbformat_minor": 5
151+
}

‎src/openlayer/lib/integrations/langchain_callback.py‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99

1010
from ..tracing import tracer
1111

12-
LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP = {"openai-chat": "OpenAI"}
13-
PROVIDER_TO_STEP_NAME = {"OpenAI": "OpenAI Chat Completion"}
12+
LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP = {"openai-chat": "OpenAI", "chat-ollama": "Ollama"}
13+
PROVIDER_TO_STEP_NAME = {"OpenAI": "OpenAI Chat Completion", "Ollama": "Ollama Chat Completion"}
1414

1515

1616
class OpenlayerHandler(BaseCallbackHandler):
@@ -45,13 +45,16 @@ def on_chat_model_start(
4545
) -> Any:
4646
"""Run when Chat Model starts running."""
4747
self.model_parameters = kwargs.get("invocation_params", {})
48+
self.metadata = kwargs.get("metadata", {})
4849

4950
provider = self.model_parameters.get("_type", None)
5051
if provider in LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP:
5152
self.provider = LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP[provider]
5253
self.model_parameters.pop("_type")
54+
self.metadata.pop("ls_provider", None)
55+
self.metadata.pop("ls_model_type", None)
5356

54-
self.model = self.model_parameters.get("model_name", None)
57+
self.model = self.model_parameters.get("model_name", None) or self.metadata.pop("ls_model_name", None)
5558
self.output = ""
5659
self.prompt = self._langchain_messages_to_prompt(messages)
5760
self.start_time = time.time()
@@ -82,17 +85,32 @@ def on_llm_end(self, response: langchain_schema.LLMResult, **kwargs: Any) -> Any
8285
self.end_time = time.time()
8386
self.latency = (self.end_time - self.start_time) * 1000
8487

85-
if response.llm_output and "token_usage" in response.llm_output:
86-
self.prompt_tokens = response.llm_output["token_usage"].get("prompt_tokens", 0)
87-
self.completion_tokens = response.llm_output["token_usage"].get("completion_tokens", 0)
88-
self.total_tokens = response.llm_output["token_usage"].get("total_tokens", 0)
88+
if self.provider == "OpenAI":
89+
self._openai_token_information(response)
90+
elif self.provider == "Ollama":
91+
self._ollama_token_information(response)
8992

9093
for generations in response.generations:
9194
for generation in generations:
9295
self.output += generation.text.replace("\n", " ")
9396

9497
self._add_to_trace()
9598

99+
def _openai_token_information(self, response: langchain_schema.LLMResult) -> None:
100+
"""Extracts OpenAI's token information."""
101+
if response.llm_output and "token_usage" in response.llm_output:
102+
self.prompt_tokens = response.llm_output["token_usage"].get("prompt_tokens", 0)
103+
self.completion_tokens = response.llm_output["token_usage"].get("completion_tokens", 0)
104+
self.total_tokens = response.llm_output["token_usage"].get("total_tokens", 0)
105+
106+
def _ollama_token_information(self, response: langchain_schema.LLMResult) -> None:
107+
"""Extracts Ollama's token information."""
108+
generation_info = response.generations[0][0].generation_info
109+
if generation_info:
110+
self.prompt_tokens = generation_info.get("prompt_eval_count", 0)
111+
self.completion_tokens = generation_info.get("eval_count", 0)
112+
self.total_tokens = self.prompt_tokens + self.completion_tokens
113+
96114
def _add_to_trace(self) -> None:
97115
"""Adds to the trace."""
98116
name = PROVIDER_TO_STEP_NAME.get(self.provider, "Chat Completion Model")
@@ -109,7 +127,7 @@ def _add_to_trace(self) -> None:
109127
model_parameters=self.model_parameters,
110128
prompt_tokens=self.prompt_tokens,
111129
completion_tokens=self.completion_tokens,
112-
metadata=self.metatada,
130+
metadata=self.metadata,
113131
)
114132

115133
def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:

0 commit comments

Comments
 (0)