Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions structllm/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,8 +111,9 @@ def parse(
# If all parsing attempts fail, return raw response
return StructuredResponse(raw_response=response)

@staticmethod
def _prepare_messages(
self, messages: List[Dict[str, str]], schema: Dict[str, Any]
messages: List[Dict[str, str]], schema: Dict[str, Any]
) -> List[Dict[str, str]]:
"""Prepare messages with structured output instructions."""

Expand Down Expand Up @@ -148,7 +149,8 @@ def _prepare_messages(

return structured_messages

def _extract_json(self, content: str) -> Optional[Dict[str, Any]]:
@staticmethod
def _extract_json(content: str) -> Optional[Dict[str, Any]]:
"""Extract JSON from content that might contain additional text."""
# Try to find JSON within the content
content = content.strip()
Expand Down
21 changes: 14 additions & 7 deletions tests/test_structllm.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,14 @@ class CalendarEvent(BaseModel):
class TestStructLLM:
"""Test cases for StructLLM client."""

def test_init_default(self):
@staticmethod
def test_init_default():
"""Test StructLLM initialization with default parameters."""
client = StructLLM()
assert client is not None

def test_init_with_params(self):
@staticmethod
def test_init_with_params():
"""Test StructLLM initialization with custom parameters."""
with patch("structllm.client.litellm") as mock_litellm:
client = StructLLM(api_base="http://localhost:1234/v1", api_key="test-key")
Expand Down Expand Up @@ -130,7 +132,8 @@ def test_parse_failure_returns_raw(self, mock_completion):
assert response.output_parsed is None
assert response.raw_response == mock_response

def test_prepare_messages_with_system(self):
@staticmethod
def test_prepare_messages_with_system():
"""Test message preparation with existing system message."""
client = StructLLM()
messages = [
Expand All @@ -147,7 +150,8 @@ def test_prepare_messages_with_system(self):
assert "JSON Schema:" in result[0]["content"]
assert result[1]["role"] == "user"

def test_prepare_messages_without_system(self):
@staticmethod
def test_prepare_messages_without_system():
"""Test message preparation without system message."""
client = StructLLM()
messages = [{"role": "user", "content": "Extract info."}]
Expand All @@ -160,7 +164,8 @@ def test_prepare_messages_without_system(self):
assert "JSON Schema:" in result[0]["content"]
assert result[1]["role"] == "user"

def test_extract_json_valid(self):
@staticmethod
def test_extract_json_valid():
"""Test JSON extraction from text."""
client = StructLLM()
content = 'Here is the result: {"name": "test", "value": 123} and more text'
Expand All @@ -171,7 +176,8 @@ def test_extract_json_valid(self):
assert result["name"] == "test"
assert result["value"] == 123

def test_extract_json_invalid(self):
@staticmethod
def test_extract_json_invalid():
"""Test JSON extraction with no valid JSON."""
client = StructLLM()
content = "This contains no JSON at all"
Expand All @@ -180,7 +186,8 @@ def test_extract_json_invalid(self):

assert result is None

def test_extract_json_malformed(self):
@staticmethod
def test_extract_json_malformed():
"""Test JSON extraction with malformed JSON."""
client = StructLLM()
content = 'Here is broken JSON: {"name": "test", "value":} incomplete'
Expand Down