Skip to content

Commit 48ee322

Browse files
author
Vadym Barda
authored
partners: add xAI chat integration (#28032)
1 parent 2898b95 commit 48ee322

File tree

26 files changed

+3378
-0
lines changed

26 files changed

+3378
-0
lines changed

‎.github/workflows/_integration_test.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ jobs:
7979
VOYAGE_API_KEY: ${{ secrets.VOYAGE_API_KEY }}
8080
COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
8181
UPSTAGE_API_KEY: ${{ secrets.UPSTAGE_API_KEY }}
82+
XAI_API_KEY: ${{ secrets.XAI_API_KEY }}
8283
run: |
8384
make integration_tests
8485

‎.github/workflows/_release.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,7 @@ jobs:
308308
VOYAGE_API_KEY: ${{ secrets.VOYAGE_API_KEY }}
309309
UPSTAGE_API_KEY: ${{ secrets.UPSTAGE_API_KEY }}
310310
FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
311+
XAI_API_KEY: ${{ secrets.XAI_API_KEY }}
311312
run: make integration_tests
312313
working-directory: ${{ inputs.working-directory }}
313314

‎docs/docs/integrations/chat/xai.ipynb

Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "raw",
5+
"id": "afaf8039",
6+
"metadata": {},
7+
"source": [
8+
"---\n",
9+
"sidebar_label: xAI\n",
10+
"---"
11+
]
12+
},
13+
{
14+
"cell_type": "markdown",
15+
"id": "e49f1e0d",
16+
"metadata": {},
17+
"source": [
18+
"# ChatXAI\n",
19+
"\n",
20+
"\n",
21+
"This page will help you get started with xAI [chat models](../../concepts/chat_models.mdx). For detailed documentation of all `ChatXAI` features and configurations head to the [API reference](https://python.langchain.com/api_reference/xai/chat_models/langchain_xai.chat_models.ChatXAI.html).\n",
22+
"\n",
23+
"[xAI](https://console.x.ai/) offers an API to interact with Grok models.\n",
24+
"\n",
25+
"## Overview\n",
26+
"### Integration details\n",
27+
"\n",
28+
"| Class | Package | Local | Serializable | [JS support](https://js.langchain.com/docs/integrations/chat/xai) | Package downloads | Package latest |\n",
29+
"| :--- | :--- | :---: | :---: | :---: | :---: | :---: |\n",
30+
"| [ChatXAI](https://python.langchain.com/api_reference/xai/chat_models/langchain_xai.chat_models.ChatXAI.html) | [langchain-xai](https://python.langchain.com/api_reference/xai/index.html) | ❌ | beta | ✅ | ![PyPI - Downloads](https://img.shields.io/pypi/dm/langchain-xai?style=flat-square&label=%20) | ![PyPI - Version](https://img.shields.io/pypi/v/langchain-xai?style=flat-square&label=%20) |\n",
31+
"\n",
32+
"### Model features\n",
33+
"| [Tool calling](../../how_to/tool_calling.ipynb) | [Structured output](../../how_to/structured_output.ipynb) | JSON mode | [Image input](../../how_to/multimodal_inputs.ipynb) | Audio input | Video input | [Token-level streaming](../../how_to/chat_streaming.ipynb) | Native async | [Token usage](../../how_to/chat_token_usage_tracking.ipynb) | [Logprobs](../../how_to/logprobs.ipynb) |\n",
34+
"| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n",
35+
"| ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | \n",
36+
"\n",
37+
"## Setup\n",
38+
"\n",
39+
"To access xAI models you'll need to create an xAI account, get an API key, and install the `langchain-xai` integration package.\n",
40+
"\n",
41+
"### Credentials\n",
42+
"\n",
43+
"Head to [this page](https://console.x.ai/) to sign up for xAI and generate an API key. Once you've done this set the `XAI_API_KEY` environment variable:"
44+
]
45+
},
46+
{
47+
"cell_type": "code",
48+
"execution_count": null,
49+
"id": "433e8d2b-9519-4b49-b2c4-7ab65b046c94",
50+
"metadata": {},
51+
"outputs": [],
52+
"source": [
53+
"import getpass\n",
54+
"import os\n",
55+
"\n",
56+
"if \"XAI_API_KEY\" not in os.environ:\n",
57+
" os.environ[\"XAI_API_KEY\"] = getpass.getpass(\"Enter your xAI API key: \")"
58+
]
59+
},
60+
{
61+
"cell_type": "markdown",
62+
"id": "72ee0c4b-9764-423a-9dbf-95129e185210",
63+
"metadata": {},
64+
"source": [
65+
"If you want to get automated tracing of your model calls you can also set your [LangSmith](https://docs.smith.langchain.com/) API key by uncommenting below:"
66+
]
67+
},
68+
{
69+
"cell_type": "code",
70+
"execution_count": 2,
71+
"id": "a15d341e-3e26-4ca3-830b-5aab30ed66de",
72+
"metadata": {},
73+
"outputs": [],
74+
"source": [
75+
"# os.environ[\"LANGSMITH_API_KEY\"] = getpass.getpass(\"Enter your LangSmith API key: \")\n",
76+
"# os.environ[\"LANGSMITH_TRACING\"] = \"true\""
77+
]
78+
},
79+
{
80+
"cell_type": "markdown",
81+
"id": "0730d6a1-c893-4840-9817-5e5251676d5d",
82+
"metadata": {},
83+
"source": [
84+
"### Installation\n",
85+
"\n",
86+
"The LangChain xAI integration lives in the `langchain-xai` package:"
87+
]
88+
},
89+
{
90+
"cell_type": "code",
91+
"execution_count": 3,
92+
"id": "652d6238-1f87-422a-b135-f5abbb8652fc",
93+
"metadata": {},
94+
"outputs": [],
95+
"source": [
96+
"%pip install -qU langchain-xai"
97+
]
98+
},
99+
{
100+
"cell_type": "markdown",
101+
"id": "a38cde65-254d-4219-a441-068766c0d4b5",
102+
"metadata": {},
103+
"source": [
104+
"## Instantiation\n",
105+
"\n",
106+
"Now we can instantiate our model object and generate chat completions:"
107+
]
108+
},
109+
{
110+
"cell_type": "code",
111+
"execution_count": 4,
112+
"id": "cb09c344-1836-4e0c-acf8-11d13ac1dbae",
113+
"metadata": {},
114+
"outputs": [],
115+
"source": [
116+
"from langchain_xai import ChatXAI\n",
117+
"\n",
118+
"llm = ChatXAI(\n",
119+
" model=\"grok-beta\",\n",
120+
" temperature=0,\n",
121+
" max_tokens=None,\n",
122+
" timeout=None,\n",
123+
" max_retries=2,\n",
124+
" # other params...\n",
125+
")"
126+
]
127+
},
128+
{
129+
"cell_type": "markdown",
130+
"id": "2b4f3e15",
131+
"metadata": {},
132+
"source": [
133+
"## Invocation"
134+
]
135+
},
136+
{
137+
"cell_type": "code",
138+
"execution_count": 5,
139+
"id": "62e0dbc3",
140+
"metadata": {
141+
"tags": []
142+
},
143+
"outputs": [
144+
{
145+
"data": {
146+
"text/plain": [
147+
"AIMessage(content=\"J'adore programmer.\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 6, 'prompt_tokens': 30, 'total_tokens': 36, 'completion_tokens_details': None, 'prompt_tokens_details': None}, 'model_name': 'grok-beta', 'system_fingerprint': 'fp_14b89b2dfc', 'finish_reason': 'stop', 'logprobs': None}, id='run-adffb7a3-e48a-4f52-b694-340d85abe5c3-0', usage_metadata={'input_tokens': 30, 'output_tokens': 6, 'total_tokens': 36, 'input_token_details': {}, 'output_token_details': {}})"
148+
]
149+
},
150+
"execution_count": 5,
151+
"metadata": {},
152+
"output_type": "execute_result"
153+
}
154+
],
155+
"source": [
156+
"messages = [\n",
157+
" (\n",
158+
" \"system\",\n",
159+
" \"You are a helpful assistant that translates English to French. Translate the user sentence.\",\n",
160+
" ),\n",
161+
" (\"human\", \"I love programming.\"),\n",
162+
"]\n",
163+
"ai_msg = llm.invoke(messages)\n",
164+
"ai_msg"
165+
]
166+
},
167+
{
168+
"cell_type": "code",
169+
"execution_count": 6,
170+
"id": "d86145b3-bfef-46e8-b227-4dda5c9c2705",
171+
"metadata": {},
172+
"outputs": [
173+
{
174+
"name": "stdout",
175+
"output_type": "stream",
176+
"text": [
177+
"J'adore programmer.\n"
178+
]
179+
}
180+
],
181+
"source": [
182+
"print(ai_msg.content)"
183+
]
184+
},
185+
{
186+
"cell_type": "markdown",
187+
"id": "18e2bfc0-7e78-4528-a73f-499ac150dca8",
188+
"metadata": {},
189+
"source": [
190+
"## Chaining\n",
191+
"\n",
192+
"We can [chain](../../how_to/sequence.ipynb) our model with a prompt template like so:"
193+
]
194+
},
195+
{
196+
"cell_type": "code",
197+
"execution_count": 7,
198+
"id": "e197d1d7-a070-4c96-9f8a-a0e86d046e0b",
199+
"metadata": {},
200+
"outputs": [
201+
{
202+
"data": {
203+
"text/plain": [
204+
"AIMessage(content='Ich liebe das Programmieren.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 7, 'prompt_tokens': 25, 'total_tokens': 32, 'completion_tokens_details': None, 'prompt_tokens_details': None}, 'model_name': 'grok-beta', 'system_fingerprint': 'fp_14b89b2dfc', 'finish_reason': 'stop', 'logprobs': None}, id='run-569fc8dc-101b-4e6d-864e-d4fa80df2b63-0', usage_metadata={'input_tokens': 25, 'output_tokens': 7, 'total_tokens': 32, 'input_token_details': {}, 'output_token_details': {}})"
205+
]
206+
},
207+
"execution_count": 7,
208+
"metadata": {},
209+
"output_type": "execute_result"
210+
}
211+
],
212+
"source": [
213+
"from langchain_core.prompts import ChatPromptTemplate\n",
214+
"\n",
215+
"prompt = ChatPromptTemplate.from_messages(\n",
216+
" [\n",
217+
" (\n",
218+
" \"system\",\n",
219+
" \"You are a helpful assistant that translates {input_language} to {output_language}.\",\n",
220+
" ),\n",
221+
" (\"human\", \"{input}\"),\n",
222+
" ]\n",
223+
")\n",
224+
"\n",
225+
"chain = prompt | llm\n",
226+
"chain.invoke(\n",
227+
" {\n",
228+
" \"input_language\": \"English\",\n",
229+
" \"output_language\": \"German\",\n",
230+
" \"input\": \"I love programming.\",\n",
231+
" }\n",
232+
")"
233+
]
234+
},
235+
{
236+
"cell_type": "markdown",
237+
"id": "e074bce1-0994-4b83-b393-ae7aa7e21750",
238+
"metadata": {},
239+
"source": [
240+
"## Tool calling\n",
241+
"\n",
242+
"ChatXAI has a [tool calling](https://docs.x.ai/docs#capabilities) (we use \"tool calling\" and \"function calling\" interchangeably here) API that lets you describe tools and their arguments, and have the model return a JSON object with a tool to invoke and the inputs to that tool. Tool-calling is extremely useful for building tool-using chains and agents, and for getting structured outputs from models more generally.\n",
243+
"\n",
244+
"### ChatXAI.bind_tools()\n",
245+
"\n",
246+
"With `ChatXAI.bind_tools`, we can easily pass in Pydantic classes, dict schemas, LangChain tools, or even functions as tools to the model. Under the hood these are converted to an OpenAI tool schemas, which looks like:\n",
247+
"```\n",
248+
"{\n",
249+
" \"name\": \"...\",\n",
250+
" \"description\": \"...\",\n",
251+
" \"parameters\": {...} # JSONSchema\n",
252+
"}\n",
253+
"```\n",
254+
"and passed in every model invocation."
255+
]
256+
},
257+
{
258+
"cell_type": "code",
259+
"execution_count": 8,
260+
"id": "c6bfe929-ec02-46bd-9d54-76350edddabc",
261+
"metadata": {},
262+
"outputs": [],
263+
"source": [
264+
"from pydantic import BaseModel, Field\n",
265+
"\n",
266+
"\n",
267+
"class GetWeather(BaseModel):\n",
268+
" \"\"\"Get the current weather in a given location\"\"\"\n",
269+
"\n",
270+
" location: str = Field(..., description=\"The city and state, e.g. San Francisco, CA\")\n",
271+
"\n",
272+
"\n",
273+
"llm_with_tools = llm.bind_tools([GetWeather])"
274+
]
275+
},
276+
{
277+
"cell_type": "code",
278+
"execution_count": 9,
279+
"id": "5265c892-d8c2-48af-aef5-adbee1647ba6",
280+
"metadata": {},
281+
"outputs": [
282+
{
283+
"data": {
284+
"text/plain": [
285+
"AIMessage(content='I am retrieving the current weather for San Francisco.', additional_kwargs={'tool_calls': [{'id': '0', 'function': {'arguments': '{\"location\":\"San Francisco, CA\"}', 'name': 'GetWeather'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 151, 'total_tokens': 162, 'completion_tokens_details': None, 'prompt_tokens_details': None}, 'model_name': 'grok-beta', 'system_fingerprint': 'fp_14b89b2dfc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-73707da7-afec-4a52-bee1-a176b0ab8585-0', tool_calls=[{'name': 'GetWeather', 'args': {'location': 'San Francisco, CA'}, 'id': '0', 'type': 'tool_call'}], usage_metadata={'input_tokens': 151, 'output_tokens': 11, 'total_tokens': 162, 'input_token_details': {}, 'output_token_details': {}})"
286+
]
287+
},
288+
"execution_count": 9,
289+
"metadata": {},
290+
"output_type": "execute_result"
291+
}
292+
],
293+
"source": [
294+
"ai_msg = llm_with_tools.invoke(\n",
295+
" \"what is the weather like in San Francisco\",\n",
296+
")\n",
297+
"ai_msg"
298+
]
299+
},
300+
{
301+
"cell_type": "markdown",
302+
"id": "3a5bb5ca-c3ae-4a58-be67-2cd18574b9a3",
303+
"metadata": {},
304+
"source": [
305+
"## API reference\n",
306+
"\n",
307+
"For detailed documentation of all `ChatXAI` features and configurations head to the API reference: https://python.langchain.com/api_reference/xai/chat_models/langchain_xai.chat_models.ChatXAI.html"
308+
]
309+
}
310+
],
311+
"metadata": {
312+
"kernelspec": {
313+
"display_name": "Python 3 (ipykernel)",
314+
"language": "python",
315+
"name": "python3"
316+
},
317+
"language_info": {
318+
"codemirror_mode": {
319+
"name": "ipython",
320+
"version": 3
321+
},
322+
"file_extension": ".py",
323+
"mimetype": "text/x-python",
324+
"name": "python",
325+
"nbconvert_exporter": "python",
326+
"pygments_lexer": "ipython3",
327+
"version": "3.11.9"
328+
}
329+
},
330+
"nbformat": 4,
331+
"nbformat_minor": 5
332+
}

0 commit comments

Comments
 (0)