dbf71a6594
Fix Zotero table models for Zotero 9:
- Remove stale Annotations/Highlights/Transaction* models
- Add ItemAnnotations, RetractedItems, DeletedCollections,
DeletedSearches, DbDebug1
- Fix ItemAttachments, Libraries, Users column mismatches
New test files covering all major modules:
- cli/{bib,prisma,rec,zot,mail,run} deep exercising tests
- mail/{droplet,postmark,resend,cloudflare} lifecycle tests
- bib/{iom,oig,pincite,sync,regulations_gov,email_ingest,format,store}
- prisma/{vpn,fetch,export,llm,screen,eligibility,extract,project,ingest,flow}
- aco/lake/{unity,quality,deploy} + api/aco coverage gaps
- zot/{ops,db,extract,duck} + rec/{report,engine,base,pricers}
- pfs/{pipe,rules,eq,files}
Add pytest-xdist for parallel test execution.
Tracks #353
246 lines
8.2 KiB
Python
246 lines
8.2 KiB
Python
"""Exercise prisma.llm — AnthropicProvider, OpenAICompatProvider, make_provider."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from prisma.llm import (
|
|
LLMCall,
|
|
LLMMessage,
|
|
LLMTool,
|
|
make_provider,
|
|
)
|
|
|
|
|
|
class TestAnthropicProvider:
|
|
def test_complete_text(self):
|
|
mock_anthropic = MagicMock()
|
|
mock_client = MagicMock()
|
|
mock_anthropic.Anthropic.return_value = mock_client
|
|
|
|
block = MagicMock()
|
|
block.type = "text"
|
|
block.text = "pong"
|
|
resp = MagicMock()
|
|
resp.content = [block]
|
|
resp.usage = MagicMock(
|
|
input_tokens=10,
|
|
output_tokens=5,
|
|
cache_creation_input_tokens=0,
|
|
cache_read_input_tokens=0,
|
|
)
|
|
mock_client.messages.create.return_value = resp
|
|
|
|
with patch.dict(sys.modules, {"anthropic": mock_anthropic}):
|
|
from prisma.llm import AnthropicProvider
|
|
|
|
provider = AnthropicProvider(api_key="test-key")
|
|
call = LLMCall(
|
|
messages=[LLMMessage(role="user", content="ping")],
|
|
max_tokens=10,
|
|
)
|
|
result = provider.complete(call)
|
|
assert result.text == "pong"
|
|
assert result.usage["input_tokens"] == 10
|
|
|
|
def test_complete_with_system_and_cache(self):
|
|
mock_anthropic = MagicMock()
|
|
mock_client = MagicMock()
|
|
mock_anthropic.Anthropic.return_value = mock_client
|
|
|
|
block = MagicMock()
|
|
block.type = "text"
|
|
block.text = "ok"
|
|
resp = MagicMock()
|
|
resp.content = [block]
|
|
resp.usage = MagicMock(
|
|
input_tokens=10,
|
|
output_tokens=1,
|
|
cache_creation_input_tokens=5,
|
|
cache_read_input_tokens=0,
|
|
)
|
|
mock_client.messages.create.return_value = resp
|
|
|
|
with patch.dict(sys.modules, {"anthropic": mock_anthropic}):
|
|
from prisma.llm import AnthropicProvider
|
|
|
|
provider = AnthropicProvider(api_key="test-key")
|
|
call = LLMCall(
|
|
messages=[
|
|
LLMMessage(role="system", content="helpful", cache=True),
|
|
LLMMessage(role="user", content="hi", cache=True),
|
|
],
|
|
)
|
|
result = provider.complete(call)
|
|
assert result.text == "ok"
|
|
kwargs = mock_client.messages.create.call_args[1]
|
|
assert "system" in kwargs
|
|
|
|
def test_complete_with_tools(self):
|
|
mock_anthropic = MagicMock()
|
|
mock_client = MagicMock()
|
|
mock_anthropic.Anthropic.return_value = mock_client
|
|
|
|
tool_block = MagicMock()
|
|
tool_block.type = "tool_use"
|
|
tool_block.name = "classify"
|
|
tool_block.input = {"decision": "include"}
|
|
resp = MagicMock()
|
|
resp.content = [tool_block]
|
|
resp.usage = MagicMock(input_tokens=20, output_tokens=10)
|
|
mock_client.messages.create.return_value = resp
|
|
|
|
with patch.dict(sys.modules, {"anthropic": mock_anthropic}):
|
|
from prisma.llm import AnthropicProvider
|
|
|
|
provider = AnthropicProvider(api_key="test-key")
|
|
tool = LLMTool(name="classify", description="c", schema={"type": "object"})
|
|
call = LLMCall(
|
|
messages=[LLMMessage(role="user", content="test")],
|
|
tools=[tool],
|
|
force_tool="classify",
|
|
)
|
|
result = provider.complete(call)
|
|
assert len(result.tool_calls) == 1
|
|
assert result.tool_calls[0]["name"] == "classify"
|
|
|
|
|
|
class TestOpenAICompatProvider:
|
|
def test_complete_text(self):
|
|
mock_openai = MagicMock()
|
|
mock_client = MagicMock()
|
|
mock_openai.OpenAI.return_value = mock_client
|
|
|
|
msg = MagicMock()
|
|
msg.content = "hello"
|
|
msg.tool_calls = None
|
|
choice = MagicMock()
|
|
choice.message = msg
|
|
resp = MagicMock()
|
|
resp.choices = [choice]
|
|
resp.usage = MagicMock(prompt_tokens=5, completion_tokens=2)
|
|
mock_client.chat.completions.create.return_value = resp
|
|
|
|
with patch.dict(sys.modules, {"openai": mock_openai}):
|
|
from prisma.llm import OpenAICompatProvider
|
|
|
|
provider = OpenAICompatProvider(
|
|
api_key="test", base_url="http://localhost:8000/v1"
|
|
)
|
|
call = LLMCall(
|
|
messages=[LLMMessage(role="user", content="hi")],
|
|
max_tokens=10,
|
|
)
|
|
result = provider.complete(call)
|
|
assert result.text == "hello"
|
|
|
|
def test_complete_with_tools_and_force(self):
|
|
mock_openai = MagicMock()
|
|
mock_client = MagicMock()
|
|
mock_openai.OpenAI.return_value = mock_client
|
|
|
|
tc = MagicMock()
|
|
tc.function.name = "classify"
|
|
tc.function.arguments = '{"decision": "exclude"}'
|
|
msg = MagicMock()
|
|
msg.content = ""
|
|
msg.tool_calls = [tc]
|
|
choice = MagicMock()
|
|
choice.message = msg
|
|
resp = MagicMock()
|
|
resp.choices = [choice]
|
|
resp.usage = MagicMock(prompt_tokens=10, completion_tokens=5)
|
|
mock_client.chat.completions.create.return_value = resp
|
|
|
|
with patch.dict(sys.modules, {"openai": mock_openai}):
|
|
from prisma.llm import OpenAICompatProvider
|
|
|
|
provider = OpenAICompatProvider(
|
|
api_key="test", base_url="http://localhost:8000/v1"
|
|
)
|
|
tool = LLMTool(name="classify", description="d", schema={"type": "object"})
|
|
call = LLMCall(
|
|
messages=[LLMMessage(role="user", content="test")],
|
|
tools=[tool],
|
|
force_tool="classify",
|
|
)
|
|
result = provider.complete(call)
|
|
assert len(result.tool_calls) == 1
|
|
assert result.tool_calls[0]["input"]["decision"] == "exclude"
|
|
|
|
def test_bad_json_args(self):
|
|
mock_openai = MagicMock()
|
|
mock_client = MagicMock()
|
|
mock_openai.OpenAI.return_value = mock_client
|
|
|
|
tc = MagicMock()
|
|
tc.function.name = "classify"
|
|
tc.function.arguments = "not json"
|
|
msg = MagicMock()
|
|
msg.content = ""
|
|
msg.tool_calls = [tc]
|
|
choice = MagicMock()
|
|
choice.message = msg
|
|
resp = MagicMock()
|
|
resp.choices = [choice]
|
|
resp.usage = MagicMock(prompt_tokens=10, completion_tokens=5)
|
|
mock_client.chat.completions.create.return_value = resp
|
|
|
|
with patch.dict(sys.modules, {"openai": mock_openai}):
|
|
from prisma.llm import OpenAICompatProvider
|
|
|
|
provider = OpenAICompatProvider(
|
|
api_key="test", base_url="http://localhost:8000/v1"
|
|
)
|
|
call = LLMCall(
|
|
messages=[LLMMessage(role="user", content="test")],
|
|
tools=[LLMTool(name="classify", description="d", schema={})],
|
|
)
|
|
result = provider.complete(call)
|
|
assert result.tool_calls[0]["input"] == {}
|
|
|
|
def test_no_usage(self):
|
|
mock_openai = MagicMock()
|
|
mock_client = MagicMock()
|
|
mock_openai.OpenAI.return_value = mock_client
|
|
|
|
msg = MagicMock()
|
|
msg.content = "ok"
|
|
msg.tool_calls = None
|
|
choice = MagicMock()
|
|
choice.message = msg
|
|
resp = MagicMock()
|
|
resp.choices = [choice]
|
|
resp.usage = None
|
|
mock_client.chat.completions.create.return_value = resp
|
|
|
|
with patch.dict(sys.modules, {"openai": mock_openai}):
|
|
from prisma.llm import OpenAICompatProvider
|
|
|
|
provider = OpenAICompatProvider(
|
|
api_key="test", base_url="http://localhost:8000/v1"
|
|
)
|
|
result = provider.complete(
|
|
LLMCall(messages=[LLMMessage(role="user", content="hi")])
|
|
)
|
|
assert result.usage["input_tokens"] == 0
|
|
|
|
|
|
class TestMakeProviderEdge:
|
|
def test_vllm(self):
|
|
mock_openai = MagicMock()
|
|
with (
|
|
patch.dict("os.environ", {"PRISMA_LLM_PROVIDER": "vllm"}),
|
|
patch.dict(sys.modules, {"openai": mock_openai}),
|
|
):
|
|
provider = make_provider()
|
|
assert provider is not None
|
|
|
|
@patch.dict("os.environ", {"PRISMA_LLM_PROVIDER": "bogus"})
|
|
def test_unknown(self):
|
|
with pytest.raises(ValueError, match="unknown"):
|
|
make_provider()
|