mirror of
https://github.com/hwchase17/langchain.git
synced 2025-05-04 14:48:07 +00:00
a0c2281540
```python
"""python scripts/update_mypy_ruff.py"""
import glob
import tomllib
from pathlib import Path
import toml
import subprocess
import re
ROOT_DIR = Path(__file__).parents[1]
def main():
for path in glob.glob(str(ROOT_DIR / "libs/**/pyproject.toml"), recursive=True):
print(path)
with open(path, "rb") as f:
pyproject = tomllib.load(f)
try:
pyproject["tool"]["poetry"]["group"]["typing"]["dependencies"]["mypy"] = (
"^1.10"
)
pyproject["tool"]["poetry"]["group"]["lint"]["dependencies"]["ruff"] = (
"^0.5"
)
except KeyError:
continue
with open(path, "w") as f:
toml.dump(pyproject, f)
cwd = "/".join(path.split("/")[:-1])
completed = subprocess.run(
"poetry lock --no-update; poetry install --with typing; poetry run mypy . --no-color",
cwd=cwd,
shell=True,
capture_output=True,
text=True,
)
logs = completed.stdout.split("\n")
to_ignore = {}
for l in logs:
if re.match("^(.*)\:(\d+)\: error:.*\[(.*)\]", l):
path, line_no, error_type = re.match(
"^(.*)\:(\d+)\: error:.*\[(.*)\]", l
).groups()
if (path, line_no) in to_ignore:
to_ignore[(path, line_no)].append(error_type)
else:
to_ignore[(path, line_no)] = [error_type]
print(len(to_ignore))
for (error_path, line_no), error_types in to_ignore.items():
all_errors = ", ".join(error_types)
full_path = f"{cwd}/{error_path}"
try:
with open(full_path, "r") as f:
file_lines = f.readlines()
except FileNotFoundError:
continue
file_lines[int(line_no) - 1] = (
file_lines[int(line_no) - 1][:-1] + f" # type: ignore[{all_errors}]\n"
)
with open(full_path, "w") as f:
f.write("".join(file_lines))
subprocess.run(
"poetry run ruff format .; poetry run ruff --select I --fix .",
cwd=cwd,
shell=True,
capture_output=True,
text=True,
)
if __name__ == "__main__":
main()
```
91 lines
2.8 KiB
Python
91 lines
2.8 KiB
Python
"""Test Anthropic Chat API wrapper."""
|
|
|
|
from typing import List
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
BaseMessage,
|
|
HumanMessage,
|
|
SystemMessage,
|
|
)
|
|
|
|
from langchain_community.chat_models import BedrockChat
|
|
from langchain_community.chat_models.meta import convert_messages_to_prompt_llama
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("messages", "expected"),
|
|
[
|
|
([HumanMessage(content="Hello")], "[INST] Hello [/INST]"),
|
|
(
|
|
[HumanMessage(content="Hello"), AIMessage(content="Answer:")],
|
|
"[INST] Hello [/INST]\nAnswer:",
|
|
),
|
|
(
|
|
[
|
|
SystemMessage(content="You're an assistant"),
|
|
HumanMessage(content="Hello"),
|
|
AIMessage(content="Answer:"),
|
|
],
|
|
"<<SYS>> You're an assistant <</SYS>>\n[INST] Hello [/INST]\nAnswer:",
|
|
),
|
|
],
|
|
)
|
|
def test_formatting(messages: List[BaseMessage], expected: str) -> None:
|
|
result = convert_messages_to_prompt_llama(messages)
|
|
assert result == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_id",
|
|
["anthropic.claude-v2", "amazon.titan-text-express-v1"],
|
|
)
|
|
def test_different_models_bedrock(model_id: str) -> None:
|
|
provider = model_id.split(".")[0]
|
|
client = MagicMock()
|
|
respbody = MagicMock()
|
|
if provider == "anthropic":
|
|
respbody.read.return_value = MagicMock(
|
|
decode=MagicMock(return_value=b'{"completion":"Hi back"}'),
|
|
)
|
|
client.invoke_model.return_value = {"body": respbody}
|
|
elif provider == "amazon":
|
|
respbody.read.return_value = '{"results": [{"outputText": "Hi back"}]}'
|
|
client.invoke_model.return_value = {"body": respbody}
|
|
|
|
model = BedrockChat(model_id=model_id, client=client)
|
|
|
|
# should not throw an error
|
|
model.invoke("hello there")
|
|
|
|
|
|
def test_bedrock_combine_llm_output() -> None:
|
|
model_id = "anthropic.claude-3-haiku-20240307-v1:0"
|
|
client = MagicMock()
|
|
llm_outputs = [
|
|
{
|
|
"model_id": "anthropic.claude-3-haiku-20240307-v1:0",
|
|
"usage": {
|
|
"completion_tokens": 1,
|
|
"prompt_tokens": 2,
|
|
"total_tokens": 3,
|
|
},
|
|
},
|
|
{
|
|
"model_id": "anthropic.claude-3-haiku-20240307-v1:0",
|
|
"usage": {
|
|
"completion_tokens": 1,
|
|
"prompt_tokens": 2,
|
|
"total_tokens": 3,
|
|
},
|
|
},
|
|
]
|
|
model = BedrockChat(model_id=model_id, client=client)
|
|
final_output = model._combine_llm_outputs(llm_outputs) # type: ignore[arg-type]
|
|
assert final_output["model_id"] == model_id
|
|
assert final_output["usage"]["completion_tokens"] == 2
|
|
assert final_output["usage"]["prompt_tokens"] == 4
|
|
assert final_output["usage"]["total_tokens"] == 6
|