Skip to content

Commit c0fa6bb

Browse files
dot-agiareibman
andauthored
[FEATURE] Add support for Mistral via the mistral Python SDK (#374)
* add mistral support * linting * fix typo * add tests * add examples notebook * linting * fix langchain typo in pyproject.toml (updated to 0.2.14) * fix mistralai import and `undo_override` function * add mistral to readme * fix typo * modified self.llm_event to llm_event * refactoring * black * rename examples directory * fix merge * init merge * updated model name so that tokencost will recognize this as a mistral model * black lint --------- Co-authored-by: reibs <areibman@gmail.com>
1 parent 1b8f473 commit c0fa6bb

6 files changed

Lines changed: 692 additions & 1 deletion

File tree

README.md

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,149 @@ async def main() -> None:
370370
print(message.content)
371371

372372

373+
await main()
374+
```
375+
</details>
376+
377+
### Mistral 〽️
378+
379+
Track agents built with the Anthropic Python SDK (>=0.32.0).
380+
381+
- [AgentOps integration example](./examples/mistral//mistral_example.ipynb)
382+
- [Official Mistral documentation](https://docs.mistral.ai)
383+
384+
<details>
385+
<summary>Installation</summary>
386+
387+
```bash
388+
pip install mistralai
389+
```
390+
391+
Sync
392+
393+
```python python
394+
from mistralai import Mistral
395+
import agentops
396+
397+
# Beginning of program's code (i.e. main.py, __init__.py)
398+
agentops.init(<INSERT YOUR API KEY HERE>)
399+
400+
client = Mistral(
401+
# This is the default and can be omitted
402+
api_key=os.environ.get("MISTRAL_API_KEY"),
403+
)
404+
405+
message = client.chat.complete(
406+
messages=[
407+
{
408+
"role": "user",
409+
"content": "Tell me a cool fact about AgentOps",
410+
}
411+
],
412+
model="open-mistral-nemo",
413+
)
414+
print(message.choices[0].message.content)
415+
416+
agentops.end_session('Success')
417+
```
418+
419+
Streaming
420+
421+
```python python
422+
from mistralai import Mistral
423+
import agentops
424+
425+
# Beginning of program's code (i.e. main.py, __init__.py)
426+
agentops.init(<INSERT YOUR API KEY HERE>)
427+
428+
client = Mistral(
429+
# This is the default and can be omitted
430+
api_key=os.environ.get("MISTRAL_API_KEY"),
431+
)
432+
433+
message = client.chat.stream(
434+
messages=[
435+
{
436+
"role": "user",
437+
"content": "Tell me something cool about streaming agents",
438+
}
439+
],
440+
model="open-mistral-nemo",
441+
)
442+
443+
response = ""
444+
for event in message:
445+
if event.data.choices[0].finish_reason == "stop":
446+
print("\n")
447+
print(response)
448+
print("\n")
449+
else:
450+
response += event.text
451+
452+
agentops.end_session('Success')
453+
```
454+
455+
Async
456+
457+
```python python
458+
import asyncio
459+
from mistralai import Mistral
460+
461+
client = Mistral(
462+
# This is the default and can be omitted
463+
api_key=os.environ.get("MISTRAL_API_KEY"),
464+
)
465+
466+
467+
async def main() -> None:
468+
message = await client.chat.complete_async(
469+
messages=[
470+
{
471+
"role": "user",
472+
"content": "Tell me something interesting about async agents",
473+
}
474+
],
475+
model="open-mistral-nemo",
476+
)
477+
print(message.choices[0].message.content)
478+
479+
480+
await main()
481+
```
482+
483+
Async Streaming
484+
485+
```python python
486+
import asyncio
487+
from mistralai import Mistral
488+
489+
client = Mistral(
490+
# This is the default and can be omitted
491+
api_key=os.environ.get("MISTRAL_API_KEY"),
492+
)
493+
494+
495+
async def main() -> None:
496+
message = await client.chat.stream_async(
497+
messages=[
498+
{
499+
"role": "user",
500+
"content": "Tell me something interesting about async streaming agents",
501+
}
502+
],
503+
model="open-mistral-nemo",
504+
)
505+
506+
response = ""
507+
async for event in message:
508+
if event.data.choices[0].finish_reason == "stop":
509+
print("\n")
510+
print(response)
511+
print("\n")
512+
else:
513+
response += event.text
514+
515+
373516
await main()
374517
```
375518
</details>

agentops/llms/__init__.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from .ollama import OllamaProvider
1414
from .openai import OpenAiProvider
1515
from .anthropic import AnthropicProvider
16+
from .mistral import MistralProvider
1617
from .ai21 import AI21Provider
1718

1819
original_func = {}
@@ -40,6 +41,9 @@ class LlmTracker:
4041
"anthropic": {
4142
"0.32.0": ("completions.create",),
4243
},
44+
"mistralai": {
45+
"1.0.1": ("chat.complete", "chat.stream"),
46+
},
4347
"ai21": {
4448
"2.0.0": (
4549
"chat.completions.create",
@@ -142,6 +146,17 @@ def override_api(self):
142146
f"Only Anthropic>=0.32.0 supported. v{module_version} found."
143147
)
144148

149+
if api == "mistralai":
150+
module_version = version(api)
151+
152+
if Version(module_version) >= parse("1.0.1"):
153+
provider = MistralProvider(self.client)
154+
provider.override()
155+
else:
156+
logger.warning(
157+
f"Only MistralAI>=1.0.1 supported. v{module_version} found."
158+
)
159+
145160
if api == "ai21":
146161
module_version = version(api)
147162

@@ -165,4 +180,5 @@ def stop_instrumenting(self):
165180
LiteLLMProvider(self.client).undo_override()
166181
OllamaProvider(self.client).undo_override()
167182
AnthropicProvider(self.client).undo_override()
183+
MistralProvider(self.client).undo_override()
168184
AI21Provider(self.client).undo_override()

0 commit comments

Comments
 (0)