ara-cli 0.1.10.5__py3-none-any.whl → 0.1.14.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ara_cli/__init__.py +51 -6
- ara_cli/__main__.py +87 -75
- ara_cli/ara_command_action.py +189 -101
- ara_cli/ara_config.py +187 -128
- ara_cli/ara_subcommands/common.py +2 -2
- ara_cli/ara_subcommands/config.py +221 -0
- ara_cli/ara_subcommands/convert.py +107 -0
- ara_cli/ara_subcommands/fetch.py +41 -0
- ara_cli/ara_subcommands/fetch_agents.py +22 -0
- ara_cli/ara_subcommands/fetch_scripts.py +19 -0
- ara_cli/ara_subcommands/fetch_templates.py +15 -10
- ara_cli/ara_subcommands/list.py +97 -23
- ara_cli/ara_subcommands/prompt.py +266 -106
- ara_cli/artefact_autofix.py +117 -64
- ara_cli/artefact_converter.py +355 -0
- ara_cli/artefact_creator.py +41 -17
- ara_cli/artefact_lister.py +3 -3
- ara_cli/artefact_models/artefact_model.py +1 -1
- ara_cli/artefact_models/artefact_templates.py +0 -9
- ara_cli/artefact_models/feature_artefact_model.py +8 -8
- ara_cli/artefact_reader.py +62 -43
- ara_cli/artefact_scan.py +39 -17
- ara_cli/chat.py +300 -71
- ara_cli/chat_agent/__init__.py +0 -0
- ara_cli/chat_agent/agent_process_manager.py +155 -0
- ara_cli/chat_script_runner/__init__.py +0 -0
- ara_cli/chat_script_runner/script_completer.py +23 -0
- ara_cli/chat_script_runner/script_finder.py +41 -0
- ara_cli/chat_script_runner/script_lister.py +36 -0
- ara_cli/chat_script_runner/script_runner.py +36 -0
- ara_cli/chat_web_search/__init__.py +0 -0
- ara_cli/chat_web_search/web_search.py +263 -0
- ara_cli/children_contribution_updater.py +737 -0
- ara_cli/classifier.py +34 -0
- ara_cli/commands/agent_run_command.py +98 -0
- ara_cli/commands/fetch_agents_command.py +106 -0
- ara_cli/commands/fetch_scripts_command.py +43 -0
- ara_cli/commands/fetch_templates_command.py +39 -0
- ara_cli/commands/fetch_templates_commands.py +39 -0
- ara_cli/commands/list_agents_command.py +39 -0
- ara_cli/commands/load_command.py +4 -3
- ara_cli/commands/load_image_command.py +1 -1
- ara_cli/commands/read_command.py +23 -27
- ara_cli/completers.py +95 -35
- ara_cli/constants.py +2 -0
- ara_cli/directory_navigator.py +37 -4
- ara_cli/error_handler.py +26 -11
- ara_cli/file_loaders/document_reader.py +0 -178
- ara_cli/file_loaders/factories/__init__.py +0 -0
- ara_cli/file_loaders/factories/document_reader_factory.py +32 -0
- ara_cli/file_loaders/factories/file_loader_factory.py +27 -0
- ara_cli/file_loaders/file_loader.py +1 -30
- ara_cli/file_loaders/loaders/__init__.py +0 -0
- ara_cli/file_loaders/{document_file_loader.py → loaders/document_file_loader.py} +1 -1
- ara_cli/file_loaders/loaders/text_file_loader.py +47 -0
- ara_cli/file_loaders/readers/__init__.py +0 -0
- ara_cli/file_loaders/readers/docx_reader.py +49 -0
- ara_cli/file_loaders/readers/excel_reader.py +27 -0
- ara_cli/file_loaders/{markdown_reader.py → readers/markdown_reader.py} +1 -1
- ara_cli/file_loaders/readers/odt_reader.py +59 -0
- ara_cli/file_loaders/readers/pdf_reader.py +54 -0
- ara_cli/file_loaders/readers/pptx_reader.py +104 -0
- ara_cli/file_loaders/tools/__init__.py +0 -0
- ara_cli/llm_utils.py +58 -0
- ara_cli/output_suppressor.py +53 -0
- ara_cli/prompt_chat.py +20 -4
- ara_cli/prompt_extractor.py +47 -32
- ara_cli/prompt_handler.py +123 -17
- ara_cli/tag_extractor.py +8 -7
- ara_cli/template_loader.py +2 -1
- ara_cli/template_manager.py +52 -21
- ara_cli/templates/global-scripts/hello_global.py +1 -0
- ara_cli/templates/prompt-modules/commands/add_scenarios_for_new_behaviour.feature_creation_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/align_feature_with_implementation_changes.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/analyze_codebase_and_plan_tasks.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/choose_best_parent_artefact.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/create_tasks_from_artefact_content.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/create_tests_for_uncovered_modules.test_generation_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/derive_features_from_video_description.feature_creation_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/describe_agent_capabilities.agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/empty.commands.md +2 -12
- ara_cli/templates/prompt-modules/commands/execute_scoped_todos_in_task.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/explain_single_file_purpose.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/extract_file_information_bullets.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/extract_general.commands.md +12 -0
- ara_cli/templates/prompt-modules/commands/extract_markdown.commands.md +11 -0
- ara_cli/templates/prompt-modules/commands/extract_python.commands.md +13 -0
- ara_cli/templates/prompt-modules/commands/feature_add_or_modifiy_specified_behavior.commands.md +36 -0
- ara_cli/templates/prompt-modules/commands/feature_generate_initial_specified_bevahior.commands.md +53 -0
- ara_cli/templates/prompt-modules/commands/fix_failing_behave_step_definitions.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/fix_failing_pytest_tests.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/general_instruction_policy.commands.md +47 -0
- ara_cli/templates/prompt-modules/commands/generate_and_fix_pytest_tests.test_generation_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/prompt_template_tech_stack_transformer.commands.md +95 -0
- ara_cli/templates/prompt-modules/commands/python_bug_fixing_code.commands.md +34 -0
- ara_cli/templates/prompt-modules/commands/python_generate_code.commands.md +27 -0
- ara_cli/templates/prompt-modules/commands/python_refactoring_code.commands.md +39 -0
- ara_cli/templates/prompt-modules/commands/python_step_definitions_generation_and_fixing.commands.md +40 -0
- ara_cli/templates/prompt-modules/commands/python_unittest_generation_and_fixing.commands.md +48 -0
- ara_cli/templates/prompt-modules/commands/suggest_next_story_child_tasks.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/summarize_or_transcribe_media.interview_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/update_feature_to_match_implementation.feature_creation_agent.commands.md +1 -0
- ara_cli/templates/prompt-modules/commands/update_user_story_with_requirements.interview_agent.commands.md +1 -0
- ara_cli/version.py +1 -1
- {ara_cli-0.1.10.5.dist-info → ara_cli-0.1.14.0.dist-info}/METADATA +49 -11
- ara_cli-0.1.14.0.dist-info/RECORD +253 -0
- {ara_cli-0.1.10.5.dist-info → ara_cli-0.1.14.0.dist-info}/WHEEL +1 -1
- tests/test_ara_command_action.py +31 -19
- tests/test_ara_config.py +177 -90
- tests/test_artefact_autofix.py +170 -97
- tests/test_artefact_autofix_integration.py +495 -0
- tests/test_artefact_converter.py +312 -0
- tests/test_artefact_extraction.py +564 -0
- tests/test_artefact_lister.py +11 -8
- tests/test_chat.py +166 -130
- tests/test_chat_givens_images.py +603 -0
- tests/test_chat_script_runner.py +454 -0
- tests/test_children_contribution_updater.py +98 -0
- tests/test_document_loader_office.py +267 -0
- tests/test_llm_utils.py +164 -0
- tests/test_prompt_chat.py +343 -0
- tests/test_prompt_extractor.py +683 -0
- tests/test_prompt_handler.py +416 -214
- tests/test_setup_default_chat_prompt_mode.py +198 -0
- tests/test_tag_extractor.py +95 -49
- tests/test_web_search.py +467 -0
- ara_cli/file_loaders/document_readers.py +0 -233
- ara_cli/file_loaders/file_loaders.py +0 -123
- ara_cli/file_loaders/text_file_loader.py +0 -187
- ara_cli/templates/prompt-modules/blueprints/complete_pytest_unittest.blueprint.md +0 -27
- ara_cli/templates/prompt-modules/blueprints/pytest_unittest_prompt.blueprint.md +0 -32
- ara_cli/templates/prompt-modules/blueprints/task_todo_list_implement_feature_BDD_way.blueprint.md +0 -30
- ara_cli/templates/prompt-modules/commands/artefact_classification.commands.md +0 -9
- ara_cli/templates/prompt-modules/commands/artefact_extension.commands.md +0 -17
- ara_cli/templates/prompt-modules/commands/artefact_formulation.commands.md +0 -14
- ara_cli/templates/prompt-modules/commands/behave_step_generation.commands.md +0 -102
- ara_cli/templates/prompt-modules/commands/code_generation_complex.commands.md +0 -20
- ara_cli/templates/prompt-modules/commands/code_generation_simple.commands.md +0 -13
- ara_cli/templates/prompt-modules/commands/error_fixing.commands.md +0 -20
- ara_cli/templates/prompt-modules/commands/feature_file_update.commands.md +0 -18
- ara_cli/templates/prompt-modules/commands/feature_formulation.commands.md +0 -43
- ara_cli/templates/prompt-modules/commands/js_code_generation_simple.commands.md +0 -13
- ara_cli/templates/prompt-modules/commands/refactoring.commands.md +0 -15
- ara_cli/templates/prompt-modules/commands/refactoring_analysis.commands.md +0 -9
- ara_cli/templates/prompt-modules/commands/reverse_engineer_feature_file.commands.md +0 -15
- ara_cli/templates/prompt-modules/commands/reverse_engineer_program_flow.commands.md +0 -19
- ara_cli-0.1.10.5.dist-info/RECORD +0 -194
- /ara_cli/file_loaders/{binary_file_loader.py → loaders/binary_file_loader.py} +0 -0
- /ara_cli/file_loaders/{image_processor.py → tools/image_processor.py} +0 -0
- {ara_cli-0.1.10.5.dist-info → ara_cli-0.1.14.0.dist-info}/entry_points.txt +0 -0
- {ara_cli-0.1.10.5.dist-info → ara_cli-0.1.14.0.dist-info}/top_level.txt +0 -0
tests/test_web_search.py
ADDED
|
@@ -0,0 +1,467 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Unit tests for chat_web_search/web_search.py
|
|
3
|
+
|
|
4
|
+
Provides full test coverage for web search functionality.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import pytest
|
|
8
|
+
from unittest.mock import patch, MagicMock
|
|
9
|
+
from ara_cli.chat_web_search.web_search import (
|
|
10
|
+
is_web_search_supported,
|
|
11
|
+
get_supported_models_message,
|
|
12
|
+
_get_raw_model_name,
|
|
13
|
+
_deduplicate_citations,
|
|
14
|
+
_format_citations,
|
|
15
|
+
_create_chunk,
|
|
16
|
+
_extract_openai_citations,
|
|
17
|
+
_extract_anthropic_text_citations,
|
|
18
|
+
_extract_anthropic_search_results,
|
|
19
|
+
perform_openai_web_search,
|
|
20
|
+
perform_anthropic_web_search,
|
|
21
|
+
perform_web_search_completion,
|
|
22
|
+
OPENAI_WEB_SEARCH_MODELS,
|
|
23
|
+
ANTHROPIC_WEB_SEARCH_MODELS,
|
|
24
|
+
)
|
|
25
|
+
from ara_cli.error_handler import AraError
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
# =============================================================================
|
|
29
|
+
# Tests for is_web_search_supported
|
|
30
|
+
# =============================================================================
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class TestIsWebSearchSupported:
|
|
34
|
+
"""Tests for is_web_search_supported function."""
|
|
35
|
+
|
|
36
|
+
@pytest.mark.parametrize(
|
|
37
|
+
"model",
|
|
38
|
+
[
|
|
39
|
+
"gpt-5",
|
|
40
|
+
"gpt-5.1",
|
|
41
|
+
"o3",
|
|
42
|
+
"o4-mini",
|
|
43
|
+
"openai/gpt-5",
|
|
44
|
+
"openai/o4-mini",
|
|
45
|
+
"gpt-5-search-api",
|
|
46
|
+
"gpt-4o-search-preview",
|
|
47
|
+
],
|
|
48
|
+
)
|
|
49
|
+
def test_openai_models_supported(self, model):
|
|
50
|
+
"""OpenAI web search models are supported."""
|
|
51
|
+
supported, provider = is_web_search_supported(model)
|
|
52
|
+
assert supported is True
|
|
53
|
+
assert provider == "openai"
|
|
54
|
+
|
|
55
|
+
@pytest.mark.parametrize(
|
|
56
|
+
"model",
|
|
57
|
+
[
|
|
58
|
+
"claude-sonnet-4-5-20250929",
|
|
59
|
+
"claude-sonnet-4-20250514",
|
|
60
|
+
"anthropic/claude-sonnet-4-5-20250929",
|
|
61
|
+
"claude-haiku-4-5-20251001",
|
|
62
|
+
"claude-opus-4-20250514",
|
|
63
|
+
],
|
|
64
|
+
)
|
|
65
|
+
def test_anthropic_models_supported(self, model):
|
|
66
|
+
"""Anthropic web search models are supported."""
|
|
67
|
+
supported, provider = is_web_search_supported(model)
|
|
68
|
+
assert supported is True
|
|
69
|
+
assert provider == "anthropic"
|
|
70
|
+
|
|
71
|
+
def test_unsupported_model_returns_false(self):
|
|
72
|
+
"""Unsupported models return False."""
|
|
73
|
+
supported, provider = is_web_search_supported("gpt-4o")
|
|
74
|
+
assert supported is False
|
|
75
|
+
assert provider is None
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
# =============================================================================
|
|
79
|
+
# Tests for get_supported_models_message
|
|
80
|
+
# =============================================================================
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
class TestGetSupportedModelsMessage:
|
|
84
|
+
"""Tests for get_supported_models_message function."""
|
|
85
|
+
|
|
86
|
+
def test_includes_model_name(self):
|
|
87
|
+
"""Message includes the unsupported model name."""
|
|
88
|
+
result = get_supported_models_message("unsupported-model")
|
|
89
|
+
assert "unsupported-model" in result
|
|
90
|
+
|
|
91
|
+
def test_lists_openai_models(self):
|
|
92
|
+
"""Message lists OpenAI models."""
|
|
93
|
+
result = get_supported_models_message("test")
|
|
94
|
+
assert "gpt-5" in result
|
|
95
|
+
assert "OpenAI" in result
|
|
96
|
+
|
|
97
|
+
def test_lists_anthropic_models(self):
|
|
98
|
+
"""Message lists Anthropic models."""
|
|
99
|
+
result = get_supported_models_message("test")
|
|
100
|
+
assert "claude" in result
|
|
101
|
+
assert "Anthropic" in result
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
# =============================================================================
|
|
105
|
+
# Tests for _get_raw_model_name
|
|
106
|
+
# =============================================================================
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
class TestGetRawModelName:
|
|
110
|
+
"""Tests for _get_raw_model_name function."""
|
|
111
|
+
|
|
112
|
+
def test_strips_openai_prefix(self):
|
|
113
|
+
"""Strips openai/ prefix."""
|
|
114
|
+
result = _get_raw_model_name("openai/gpt-5")
|
|
115
|
+
assert result == "gpt-5"
|
|
116
|
+
|
|
117
|
+
def test_strips_anthropic_prefix(self):
|
|
118
|
+
"""Strips anthropic/ prefix."""
|
|
119
|
+
result = _get_raw_model_name("anthropic/claude-sonnet-4-20250514")
|
|
120
|
+
assert result == "claude-sonnet-4-20250514"
|
|
121
|
+
|
|
122
|
+
def test_returns_unchanged_without_prefix(self):
|
|
123
|
+
"""Returns model unchanged when no prefix."""
|
|
124
|
+
result = _get_raw_model_name("gpt-5")
|
|
125
|
+
assert result == "gpt-5"
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
# =============================================================================
|
|
129
|
+
# Tests for _deduplicate_citations
|
|
130
|
+
# =============================================================================
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
class TestDeduplicateCitations:
|
|
134
|
+
"""Tests for _deduplicate_citations function."""
|
|
135
|
+
|
|
136
|
+
def test_removes_duplicate_urls(self):
|
|
137
|
+
"""Removes citations with duplicate URLs."""
|
|
138
|
+
citations = [
|
|
139
|
+
{"title": "First", "url": "https://example.com"},
|
|
140
|
+
{"title": "Second", "url": "https://example.com"},
|
|
141
|
+
{"title": "Third", "url": "https://other.com"},
|
|
142
|
+
]
|
|
143
|
+
result = _deduplicate_citations(citations)
|
|
144
|
+
assert len(result) == 2
|
|
145
|
+
assert result[0]["title"] == "First"
|
|
146
|
+
assert result[1]["title"] == "Third"
|
|
147
|
+
|
|
148
|
+
def test_preserves_order(self):
|
|
149
|
+
"""Preserves order of first occurrences."""
|
|
150
|
+
citations = [
|
|
151
|
+
{"title": "A", "url": "https://a.com"},
|
|
152
|
+
{"title": "B", "url": "https://b.com"},
|
|
153
|
+
{"title": "C", "url": "https://c.com"},
|
|
154
|
+
]
|
|
155
|
+
result = _deduplicate_citations(citations)
|
|
156
|
+
assert [c["title"] for c in result] == ["A", "B", "C"]
|
|
157
|
+
|
|
158
|
+
def test_handles_empty_list(self):
|
|
159
|
+
"""Handles empty citation list."""
|
|
160
|
+
result = _deduplicate_citations([])
|
|
161
|
+
assert result == []
|
|
162
|
+
|
|
163
|
+
def test_skips_empty_urls(self):
|
|
164
|
+
"""Skips citations with empty URLs."""
|
|
165
|
+
citations = [
|
|
166
|
+
{"title": "First", "url": ""},
|
|
167
|
+
{"title": "Second", "url": "https://example.com"},
|
|
168
|
+
]
|
|
169
|
+
result = _deduplicate_citations(citations)
|
|
170
|
+
assert len(result) == 1
|
|
171
|
+
assert result[0]["title"] == "Second"
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
# =============================================================================
|
|
175
|
+
# Tests for _format_citations
|
|
176
|
+
# =============================================================================
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
class TestFormatCitations:
|
|
180
|
+
"""Tests for _format_citations function."""
|
|
181
|
+
|
|
182
|
+
def test_formats_markdown_links(self):
|
|
183
|
+
"""Formats citations as markdown links."""
|
|
184
|
+
citations = [
|
|
185
|
+
{"title": "Example Site", "url": "https://example.com"},
|
|
186
|
+
]
|
|
187
|
+
result = _format_citations(citations)
|
|
188
|
+
assert "[Example Site](https://example.com)" in result
|
|
189
|
+
|
|
190
|
+
def test_includes_sources_header(self):
|
|
191
|
+
"""Includes Sources header."""
|
|
192
|
+
citations = [{"title": "Test", "url": "https://test.com"}]
|
|
193
|
+
result = _format_citations(citations)
|
|
194
|
+
assert "**Sources:**" in result
|
|
195
|
+
|
|
196
|
+
def test_numbers_citations(self):
|
|
197
|
+
"""Numbers each citation."""
|
|
198
|
+
citations = [
|
|
199
|
+
{"title": "First", "url": "https://first.com"},
|
|
200
|
+
{"title": "Second", "url": "https://second.com"},
|
|
201
|
+
]
|
|
202
|
+
result = _format_citations(citations)
|
|
203
|
+
assert "1. [First]" in result
|
|
204
|
+
assert "2. [Second]" in result
|
|
205
|
+
|
|
206
|
+
def test_returns_empty_for_no_citations(self):
|
|
207
|
+
"""Returns empty string when no citations."""
|
|
208
|
+
result = _format_citations([])
|
|
209
|
+
assert result == ""
|
|
210
|
+
|
|
211
|
+
def test_handles_missing_url(self):
|
|
212
|
+
"""Handles citations without URL."""
|
|
213
|
+
citations = [{"title": "No URL", "url": ""}]
|
|
214
|
+
result = _format_citations(citations)
|
|
215
|
+
assert result == "" # Empty URL is filtered by deduplicate
|
|
216
|
+
|
|
217
|
+
|
|
218
|
+
# =============================================================================
|
|
219
|
+
# Tests for _create_chunk
|
|
220
|
+
# =============================================================================
|
|
221
|
+
|
|
222
|
+
|
|
223
|
+
class TestCreateChunk:
|
|
224
|
+
"""Tests for _create_chunk function."""
|
|
225
|
+
|
|
226
|
+
def test_creates_mock_chunk_with_content(self):
|
|
227
|
+
"""Creates mock chunk with content accessible via choices."""
|
|
228
|
+
chunk = _create_chunk("test content")
|
|
229
|
+
assert chunk.choices[0].delta.content == "test content"
|
|
230
|
+
|
|
231
|
+
|
|
232
|
+
# =============================================================================
|
|
233
|
+
# Tests for _extract_openai_citations
|
|
234
|
+
# =============================================================================
|
|
235
|
+
|
|
236
|
+
|
|
237
|
+
class TestExtractOpenaiCitations:
|
|
238
|
+
"""Tests for _extract_openai_citations function."""
|
|
239
|
+
|
|
240
|
+
def test_extracts_url_citations(self):
|
|
241
|
+
"""Extracts URL citations from response."""
|
|
242
|
+
mock_annotation = MagicMock()
|
|
243
|
+
mock_annotation.type = "url_citation"
|
|
244
|
+
mock_annotation.title = "Test Title"
|
|
245
|
+
mock_annotation.url = "https://test.com"
|
|
246
|
+
|
|
247
|
+
mock_content_item = MagicMock()
|
|
248
|
+
mock_content_item.annotations = [mock_annotation]
|
|
249
|
+
|
|
250
|
+
mock_output_item = MagicMock()
|
|
251
|
+
mock_output_item.type = "message"
|
|
252
|
+
mock_output_item.content = [mock_content_item]
|
|
253
|
+
|
|
254
|
+
mock_response = MagicMock()
|
|
255
|
+
mock_response.output = [mock_output_item]
|
|
256
|
+
|
|
257
|
+
result = _extract_openai_citations(mock_response)
|
|
258
|
+
|
|
259
|
+
assert len(result) == 1
|
|
260
|
+
assert result[0]["title"] == "Test Title"
|
|
261
|
+
assert result[0]["url"] == "https://test.com"
|
|
262
|
+
|
|
263
|
+
def test_returns_empty_for_no_output(self):
|
|
264
|
+
"""Returns empty list when no output."""
|
|
265
|
+
mock_response = MagicMock()
|
|
266
|
+
mock_response.output = None
|
|
267
|
+
|
|
268
|
+
result = _extract_openai_citations(mock_response)
|
|
269
|
+
|
|
270
|
+
assert result == []
|
|
271
|
+
|
|
272
|
+
|
|
273
|
+
# =============================================================================
|
|
274
|
+
# Tests for _extract_anthropic_text_citations
|
|
275
|
+
# =============================================================================
|
|
276
|
+
|
|
277
|
+
|
|
278
|
+
class TestExtractAnthropicTextCitations:
|
|
279
|
+
"""Tests for _extract_anthropic_text_citations function."""
|
|
280
|
+
|
|
281
|
+
def test_extracts_citations_with_url(self):
|
|
282
|
+
"""Extracts citations that have URL attribute."""
|
|
283
|
+
mock_citation = MagicMock()
|
|
284
|
+
mock_citation.url = "https://example.com"
|
|
285
|
+
mock_citation.title = "Example"
|
|
286
|
+
|
|
287
|
+
mock_content_block = MagicMock()
|
|
288
|
+
mock_content_block.citations = [mock_citation]
|
|
289
|
+
|
|
290
|
+
result = _extract_anthropic_text_citations(mock_content_block)
|
|
291
|
+
|
|
292
|
+
assert len(result) == 1
|
|
293
|
+
assert result[0]["url"] == "https://example.com"
|
|
294
|
+
|
|
295
|
+
def test_returns_empty_for_no_citations(self):
|
|
296
|
+
"""Returns empty list when no citations."""
|
|
297
|
+
mock_content_block = MagicMock()
|
|
298
|
+
mock_content_block.citations = None
|
|
299
|
+
|
|
300
|
+
result = _extract_anthropic_text_citations(mock_content_block)
|
|
301
|
+
|
|
302
|
+
assert result == []
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
# =============================================================================
|
|
306
|
+
# Tests for _extract_anthropic_search_results
|
|
307
|
+
# =============================================================================
|
|
308
|
+
|
|
309
|
+
|
|
310
|
+
class TestExtractAnthropicSearchResults:
|
|
311
|
+
"""Tests for _extract_anthropic_search_results function."""
|
|
312
|
+
|
|
313
|
+
def test_extracts_web_search_results(self):
|
|
314
|
+
"""Extracts web search result citations."""
|
|
315
|
+
mock_result = MagicMock()
|
|
316
|
+
mock_result.type = "web_search_result"
|
|
317
|
+
mock_result.title = "Search Result"
|
|
318
|
+
mock_result.url = "https://search.com"
|
|
319
|
+
|
|
320
|
+
mock_content_block = MagicMock()
|
|
321
|
+
mock_content_block.content = [mock_result]
|
|
322
|
+
|
|
323
|
+
result = _extract_anthropic_search_results(mock_content_block)
|
|
324
|
+
|
|
325
|
+
assert len(result) == 1
|
|
326
|
+
assert result[0]["title"] == "Search Result"
|
|
327
|
+
assert result[0]["url"] == "https://search.com"
|
|
328
|
+
|
|
329
|
+
def test_skips_non_search_results(self):
|
|
330
|
+
"""Skips items that aren't web_search_result type."""
|
|
331
|
+
mock_result = MagicMock()
|
|
332
|
+
mock_result.type = "other_type"
|
|
333
|
+
|
|
334
|
+
mock_content_block = MagicMock()
|
|
335
|
+
mock_content_block.content = [mock_result]
|
|
336
|
+
|
|
337
|
+
result = _extract_anthropic_search_results(mock_content_block)
|
|
338
|
+
|
|
339
|
+
assert result == []
|
|
340
|
+
|
|
341
|
+
|
|
342
|
+
# =============================================================================
|
|
343
|
+
# Tests for perform_openai_web_search
|
|
344
|
+
# =============================================================================
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
class TestPerformOpenaiWebSearch:
|
|
348
|
+
"""Tests for perform_openai_web_search function."""
|
|
349
|
+
|
|
350
|
+
@patch("openai.OpenAI")
|
|
351
|
+
@patch("os.getenv", return_value="test-api-key")
|
|
352
|
+
def test_uses_chat_completions_for_search_models(
|
|
353
|
+
self, mock_getenv, mock_openai_class
|
|
354
|
+
):
|
|
355
|
+
"""Uses Chat Completions API for search models."""
|
|
356
|
+
mock_client = MagicMock()
|
|
357
|
+
mock_openai_class.return_value = mock_client
|
|
358
|
+
|
|
359
|
+
mock_chunk = MagicMock()
|
|
360
|
+
mock_chunk.choices = [MagicMock(delta=MagicMock(content="response"))]
|
|
361
|
+
mock_client.chat.completions.create.return_value = [mock_chunk]
|
|
362
|
+
|
|
363
|
+
results = list(perform_openai_web_search("test query", "gpt-5-search-api"))
|
|
364
|
+
|
|
365
|
+
mock_client.chat.completions.create.assert_called_once()
|
|
366
|
+
|
|
367
|
+
@patch("openai.OpenAI")
|
|
368
|
+
@patch("os.getenv", return_value="test-api-key")
|
|
369
|
+
def test_uses_responses_api_for_other_models(self, mock_getenv, mock_openai_class):
|
|
370
|
+
"""Uses Responses API for non-search models."""
|
|
371
|
+
mock_client = MagicMock()
|
|
372
|
+
mock_openai_class.return_value = mock_client
|
|
373
|
+
|
|
374
|
+
mock_response = MagicMock()
|
|
375
|
+
mock_response.output_text = "response text"
|
|
376
|
+
mock_response.output = []
|
|
377
|
+
mock_client.responses.create.return_value = mock_response
|
|
378
|
+
|
|
379
|
+
results = list(perform_openai_web_search("test query", "gpt-5"))
|
|
380
|
+
|
|
381
|
+
mock_client.responses.create.assert_called_once()
|
|
382
|
+
|
|
383
|
+
|
|
384
|
+
# =============================================================================
|
|
385
|
+
# Tests for perform_anthropic_web_search
|
|
386
|
+
# =============================================================================
|
|
387
|
+
|
|
388
|
+
|
|
389
|
+
class TestPerformAnthropicWebSearch:
|
|
390
|
+
"""Tests for perform_anthropic_web_search function."""
|
|
391
|
+
|
|
392
|
+
@patch("anthropic.Anthropic")
|
|
393
|
+
@patch("os.getenv", return_value="test-api-key")
|
|
394
|
+
def test_creates_message_with_web_search_tool(
|
|
395
|
+
self, mock_getenv, mock_anthropic_class
|
|
396
|
+
):
|
|
397
|
+
"""Creates message with web_search tool."""
|
|
398
|
+
mock_client = MagicMock()
|
|
399
|
+
mock_anthropic_class.return_value = mock_client
|
|
400
|
+
|
|
401
|
+
mock_text_block = MagicMock()
|
|
402
|
+
mock_text_block.type = "text"
|
|
403
|
+
mock_text_block.text = "response text"
|
|
404
|
+
mock_text_block.citations = None
|
|
405
|
+
|
|
406
|
+
mock_response = MagicMock()
|
|
407
|
+
mock_response.content = [mock_text_block]
|
|
408
|
+
mock_client.messages.create.return_value = mock_response
|
|
409
|
+
|
|
410
|
+
results = list(
|
|
411
|
+
perform_anthropic_web_search("test query", "claude-sonnet-4-20250514")
|
|
412
|
+
)
|
|
413
|
+
|
|
414
|
+
mock_client.messages.create.assert_called_once()
|
|
415
|
+
call_args = mock_client.messages.create.call_args
|
|
416
|
+
assert any("web_search" in str(arg) for arg in call_args)
|
|
417
|
+
|
|
418
|
+
|
|
419
|
+
# =============================================================================
|
|
420
|
+
# Tests for perform_web_search_completion
|
|
421
|
+
# =============================================================================
|
|
422
|
+
|
|
423
|
+
|
|
424
|
+
class TestPerformWebSearchCompletion:
|
|
425
|
+
"""Tests for perform_web_search_completion function."""
|
|
426
|
+
|
|
427
|
+
@patch("ara_cli.chat_web_search.web_search.LLMSingleton")
|
|
428
|
+
def test_raises_for_unsupported_model(self, mock_singleton):
|
|
429
|
+
"""Raises AraError for unsupported model."""
|
|
430
|
+
mock_instance = MagicMock()
|
|
431
|
+
mock_instance.get_config_by_purpose.return_value = {"model": "gpt-4o"}
|
|
432
|
+
mock_singleton.get_instance.return_value = mock_instance
|
|
433
|
+
|
|
434
|
+
with pytest.raises(AraError) as exc_info:
|
|
435
|
+
list(perform_web_search_completion("test query"))
|
|
436
|
+
|
|
437
|
+
assert "not supported" in str(exc_info.value)
|
|
438
|
+
|
|
439
|
+
@patch("ara_cli.chat_web_search.web_search.perform_openai_web_search")
|
|
440
|
+
@patch("ara_cli.chat_web_search.web_search.LLMSingleton")
|
|
441
|
+
def test_uses_openai_for_openai_models(self, mock_singleton, mock_openai_search):
|
|
442
|
+
"""Uses OpenAI search for OpenAI models."""
|
|
443
|
+
mock_instance = MagicMock()
|
|
444
|
+
mock_instance.get_config_by_purpose.return_value = {"model": "gpt-5"}
|
|
445
|
+
mock_singleton.get_instance.return_value = mock_instance
|
|
446
|
+
mock_openai_search.return_value = iter([])
|
|
447
|
+
|
|
448
|
+
list(perform_web_search_completion("test query"))
|
|
449
|
+
|
|
450
|
+
mock_openai_search.assert_called_once()
|
|
451
|
+
|
|
452
|
+
@patch("ara_cli.chat_web_search.web_search.perform_anthropic_web_search")
|
|
453
|
+
@patch("ara_cli.chat_web_search.web_search.LLMSingleton")
|
|
454
|
+
def test_uses_anthropic_for_anthropic_models(
|
|
455
|
+
self, mock_singleton, mock_anthropic_search
|
|
456
|
+
):
|
|
457
|
+
"""Uses Anthropic search for Anthropic models."""
|
|
458
|
+
mock_instance = MagicMock()
|
|
459
|
+
mock_instance.get_config_by_purpose.return_value = {
|
|
460
|
+
"model": "claude-sonnet-4-20250514"
|
|
461
|
+
}
|
|
462
|
+
mock_singleton.get_instance.return_value = mock_instance
|
|
463
|
+
mock_anthropic_search.return_value = iter([])
|
|
464
|
+
|
|
465
|
+
list(perform_web_search_completion("test query"))
|
|
466
|
+
|
|
467
|
+
mock_anthropic_search.assert_called_once()
|
|
@@ -1,233 +0,0 @@
|
|
|
1
|
-
import os
|
|
2
|
-
from abc import ABC, abstractmethod
|
|
3
|
-
from typing import Tuple, Optional
|
|
4
|
-
|
|
5
|
-
|
|
6
|
-
class DocumentReader(ABC):
|
|
7
|
-
"""Abstract base class for document readers"""
|
|
8
|
-
|
|
9
|
-
def __init__(self, file_path: str):
|
|
10
|
-
self.file_path = file_path
|
|
11
|
-
self.base_dir = os.path.dirname(file_path)
|
|
12
|
-
|
|
13
|
-
@abstractmethod
|
|
14
|
-
def read(self, extract_images: bool = False) -> str:
|
|
15
|
-
"""Read document and optionally extract images"""
|
|
16
|
-
pass
|
|
17
|
-
|
|
18
|
-
def create_image_data_dir(self, extension_suffix: str) -> str:
|
|
19
|
-
"""
|
|
20
|
-
Create data directory for images with file extension suffix to avoid conflicts.
|
|
21
|
-
|
|
22
|
-
Returns:
|
|
23
|
-
str: Path to images directory
|
|
24
|
-
"""
|
|
25
|
-
file_name_with_ext = os.path.splitext(os.path.basename(self.file_path))[0] + f"_{extension_suffix}"
|
|
26
|
-
data_dir = os.path.join(self.base_dir, f"{file_name_with_ext}.data")
|
|
27
|
-
images_dir = os.path.join(data_dir, "images")
|
|
28
|
-
if not os.path.exists(images_dir):
|
|
29
|
-
os.makedirs(images_dir)
|
|
30
|
-
return images_dir
|
|
31
|
-
|
|
32
|
-
def save_and_describe_image(self, image_data: bytes, image_format: str,
|
|
33
|
-
save_dir: str, image_counter: int) -> Tuple[str, str]:
|
|
34
|
-
"""
|
|
35
|
-
Save image data and get its description from LLM.
|
|
36
|
-
|
|
37
|
-
Returns:
|
|
38
|
-
tuple: (relative_image_path, description)
|
|
39
|
-
"""
|
|
40
|
-
from ara_cli.prompt_handler import describe_image
|
|
41
|
-
|
|
42
|
-
# Save image
|
|
43
|
-
image_filename = f"{image_counter}.{image_format}"
|
|
44
|
-
image_path = os.path.join(save_dir, image_filename)
|
|
45
|
-
|
|
46
|
-
with open(image_path, "wb") as image_file:
|
|
47
|
-
image_file.write(image_data)
|
|
48
|
-
|
|
49
|
-
# Get image description from LLM
|
|
50
|
-
description = describe_image(image_path)
|
|
51
|
-
|
|
52
|
-
# Get relative path
|
|
53
|
-
relative_image_path = os.path.relpath(image_path, self.base_dir)
|
|
54
|
-
|
|
55
|
-
return relative_image_path, description
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
class DocxReader(DocumentReader):
|
|
59
|
-
"""Reader for DOCX files"""
|
|
60
|
-
|
|
61
|
-
def read(self, extract_images: bool = False) -> str:
|
|
62
|
-
import docx
|
|
63
|
-
|
|
64
|
-
doc = docx.Document(self.file_path)
|
|
65
|
-
text_content = '\n'.join(para.text for para in doc.paragraphs)
|
|
66
|
-
|
|
67
|
-
if not extract_images:
|
|
68
|
-
return text_content
|
|
69
|
-
|
|
70
|
-
from PIL import Image
|
|
71
|
-
import io
|
|
72
|
-
|
|
73
|
-
# Create data directory for images
|
|
74
|
-
images_dir = self.create_image_data_dir("docx")
|
|
75
|
-
|
|
76
|
-
# Extract and process images
|
|
77
|
-
image_descriptions = []
|
|
78
|
-
image_counter = 1
|
|
79
|
-
|
|
80
|
-
for rel in doc.part.rels.values():
|
|
81
|
-
if "image" in rel.reltype:
|
|
82
|
-
image_data = rel.target_part.blob
|
|
83
|
-
|
|
84
|
-
# Determine image format
|
|
85
|
-
image = Image.open(io.BytesIO(image_data))
|
|
86
|
-
image_format = image.format.lower()
|
|
87
|
-
|
|
88
|
-
# Save and describe image
|
|
89
|
-
relative_path, description = self.save_and_describe_image(
|
|
90
|
-
image_data, image_format, images_dir, image_counter
|
|
91
|
-
)
|
|
92
|
-
|
|
93
|
-
# Add formatted description to list
|
|
94
|
-
image_description = f"\nImage: {relative_path}\n[{description}]\n"
|
|
95
|
-
image_descriptions.append(image_description)
|
|
96
|
-
|
|
97
|
-
image_counter += 1
|
|
98
|
-
|
|
99
|
-
# Combine text content with image descriptions
|
|
100
|
-
if image_descriptions:
|
|
101
|
-
text_content += "\n\n### Extracted Images\n" + "\n".join(image_descriptions)
|
|
102
|
-
|
|
103
|
-
return text_content
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
class PdfReader(DocumentReader):
|
|
107
|
-
"""Reader for PDF files"""
|
|
108
|
-
|
|
109
|
-
def read(self, extract_images: bool = False) -> str:
|
|
110
|
-
import pymupdf4llm
|
|
111
|
-
|
|
112
|
-
if not extract_images:
|
|
113
|
-
return pymupdf4llm.to_markdown(self.file_path, write_images=False)
|
|
114
|
-
|
|
115
|
-
import fitz # PyMuPDF
|
|
116
|
-
|
|
117
|
-
# Create images directory
|
|
118
|
-
images_dir = self.create_image_data_dir("pdf")
|
|
119
|
-
|
|
120
|
-
# Extract text without images first
|
|
121
|
-
text_content = pymupdf4llm.to_markdown(self.file_path, write_images=False)
|
|
122
|
-
|
|
123
|
-
# Extract and process images
|
|
124
|
-
doc = fitz.open(self.file_path)
|
|
125
|
-
image_descriptions = []
|
|
126
|
-
image_counter = 1
|
|
127
|
-
|
|
128
|
-
for page_num, page in enumerate(doc):
|
|
129
|
-
image_list = page.get_images()
|
|
130
|
-
|
|
131
|
-
for img_index, img in enumerate(image_list):
|
|
132
|
-
# Extract image
|
|
133
|
-
xref = img[0]
|
|
134
|
-
base_image = doc.extract_image(xref)
|
|
135
|
-
image_bytes = base_image["image"]
|
|
136
|
-
image_ext = base_image["ext"]
|
|
137
|
-
|
|
138
|
-
# Save and describe image
|
|
139
|
-
relative_path, description = self.save_and_describe_image(
|
|
140
|
-
image_bytes, image_ext, images_dir, image_counter
|
|
141
|
-
)
|
|
142
|
-
|
|
143
|
-
# Add formatted description to list
|
|
144
|
-
image_description = f"\nImage: {relative_path}\n[{description}]\n"
|
|
145
|
-
image_descriptions.append(image_description)
|
|
146
|
-
|
|
147
|
-
image_counter += 1
|
|
148
|
-
|
|
149
|
-
doc.close()
|
|
150
|
-
|
|
151
|
-
# Combine text content with image descriptions
|
|
152
|
-
if image_descriptions:
|
|
153
|
-
text_content += "\n\n### Extracted Images\n" + "\n".join(image_descriptions)
|
|
154
|
-
|
|
155
|
-
return text_content
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
class OdtReader(DocumentReader):
|
|
159
|
-
"""Reader for ODT files"""
|
|
160
|
-
|
|
161
|
-
def read(self, extract_images: bool = False) -> str:
|
|
162
|
-
import pymupdf4llm
|
|
163
|
-
|
|
164
|
-
if not extract_images:
|
|
165
|
-
return pymupdf4llm.to_markdown(self.file_path, write_images=False)
|
|
166
|
-
|
|
167
|
-
import zipfile
|
|
168
|
-
from PIL import Image
|
|
169
|
-
import io
|
|
170
|
-
|
|
171
|
-
# Create data directory for images
|
|
172
|
-
images_dir = self.create_image_data_dir("odt")
|
|
173
|
-
|
|
174
|
-
# Get text content
|
|
175
|
-
text_content = pymupdf4llm.to_markdown(self.file_path, write_images=False)
|
|
176
|
-
|
|
177
|
-
# Extract and process images from ODT
|
|
178
|
-
image_descriptions = []
|
|
179
|
-
image_counter = 1
|
|
180
|
-
|
|
181
|
-
try:
|
|
182
|
-
with zipfile.ZipFile(self.file_path, 'r') as odt_zip:
|
|
183
|
-
# List all files in the Pictures directory
|
|
184
|
-
picture_files = [f for f in odt_zip.namelist() if f.startswith('Pictures/')]
|
|
185
|
-
|
|
186
|
-
for picture_file in picture_files:
|
|
187
|
-
# Extract image data
|
|
188
|
-
image_data = odt_zip.read(picture_file)
|
|
189
|
-
|
|
190
|
-
# Determine image format
|
|
191
|
-
image = Image.open(io.BytesIO(image_data))
|
|
192
|
-
image_format = image.format.lower()
|
|
193
|
-
|
|
194
|
-
# Save and describe image
|
|
195
|
-
relative_path, description = self.save_and_describe_image(
|
|
196
|
-
image_data, image_format, images_dir, image_counter
|
|
197
|
-
)
|
|
198
|
-
|
|
199
|
-
# Add formatted description to list
|
|
200
|
-
image_description = f"\nImage: {relative_path}\n[{description}]\n"
|
|
201
|
-
image_descriptions.append(image_description)
|
|
202
|
-
|
|
203
|
-
image_counter += 1
|
|
204
|
-
except Exception as e:
|
|
205
|
-
print(f"Warning: Could not extract images from ODT: {e}")
|
|
206
|
-
|
|
207
|
-
# Combine text content with image descriptions
|
|
208
|
-
if image_descriptions:
|
|
209
|
-
text_content += "\n\n### Extracted Images\n" + "\n".join(image_descriptions)
|
|
210
|
-
|
|
211
|
-
return text_content
|
|
212
|
-
|
|
213
|
-
|
|
214
|
-
class DocumentReaderFactory:
|
|
215
|
-
"""Factory for creating appropriate document readers"""
|
|
216
|
-
|
|
217
|
-
@staticmethod
|
|
218
|
-
def create_reader(file_path: str) -> Optional[DocumentReader]:
|
|
219
|
-
"""Create appropriate reader based on file extension"""
|
|
220
|
-
_, ext = os.path.splitext(file_path)
|
|
221
|
-
ext = ext.lower()
|
|
222
|
-
|
|
223
|
-
readers = {
|
|
224
|
-
'.docx': DocxReader,
|
|
225
|
-
'.pdf': PdfReader,
|
|
226
|
-
'.odt': OdtReader
|
|
227
|
-
}
|
|
228
|
-
|
|
229
|
-
reader_class = readers.get(ext)
|
|
230
|
-
if reader_class:
|
|
231
|
-
return reader_class(file_path)
|
|
232
|
-
|
|
233
|
-
return None
|