mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 17:31:00 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2799304270 |
No files matched your search
@@ -180,6 +180,18 @@ export OPENAI_API_KEY=xxxxxxxxxxxxxxxxxxxxxx
|
||||
praisonai --auto create a movie script about Robots in Mars
|
||||
```
|
||||
|
||||
### Query Rewriting (works with any command):
|
||||
```bash
|
||||
# Rewrite query for better results (uses QueryRewriterAgent)
|
||||
praisonai "AI trends" --query-rewrite
|
||||
|
||||
# Rewrite with search tools (agent decides when to search)
|
||||
praisonai "latest developments" --query-rewrite --rewrite-tools "internet_search"
|
||||
|
||||
# Works with any prompt
|
||||
praisonai "explain quantum computing" --query-rewrite -v
|
||||
```
|
||||
|
||||
### Deep Research CLI:
|
||||
```bash
|
||||
# Default: OpenAI (o4-mini-deep-research)
|
||||
@@ -188,10 +200,10 @@ praisonai research "What are the latest AI trends in 2025?"
|
||||
# Use Gemini
|
||||
praisonai research --model deep-research-pro "Your research query"
|
||||
|
||||
# Rewrite query for better results (uses QueryRewriterAgent)
|
||||
# Rewrite query before research
|
||||
praisonai research --query-rewrite "AI trends"
|
||||
|
||||
# Rewrite with search tools (searches first, then rewrites based on results)
|
||||
# Rewrite with search tools
|
||||
praisonai research --query-rewrite --rewrite-tools "internet_search" "AI trends"
|
||||
|
||||
# Use custom tools from file (gathers context before deep research)
|
||||
|
||||
+1
-1
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
flask \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.2.97" \
|
||||
"praisonai[api]" \
|
||||
gunicorn \
|
||||
markdown
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.2.97" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.2.97" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.2.96" \
|
||||
"praisonai>=2.2.97" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.0.167"
|
||||
version = "0.0.168"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
requires-python = ">=3.10"
|
||||
authors = [
|
||||
|
||||
Generated
+1
-1
@@ -2422,7 +2422,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.0.166"
|
||||
version = "0.0.168"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -3,8 +3,8 @@ class Praisonai < Formula
|
||||
|
||||
desc "AI tools for various AI applications"
|
||||
homepage "https://github.com/MervinPraison/PraisonAI"
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.96.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.96.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.97.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.97.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -514,9 +514,8 @@ class PraisonAI:
|
||||
in_test_env = (
|
||||
'pytest' in sys.argv[0] or
|
||||
'unittest' in sys.argv[0] or
|
||||
any('test' in arg for arg in sys.argv[1:3]) or # Check first few args for test indicators
|
||||
'pytest' in sys.modules or
|
||||
'unittest' in sys.modules
|
||||
any(arg.startswith('tests.') or arg.startswith('tests/') for arg in sys.argv[1:3]) or # Check for test module paths
|
||||
'PYTEST_CURRENT_TEST' in os.environ
|
||||
)
|
||||
|
||||
# Check if we're being used as a library (not from praisonai CLI)
|
||||
@@ -578,7 +577,7 @@ class PraisonAI:
|
||||
parser.add_argument("--goal", type=str, help="Goal for context engineering")
|
||||
parser.add_argument("--auto-analyze", action="store_true", help="Enable automatic analysis in context engineering")
|
||||
parser.add_argument("--research", action="store_true", help="Run deep research on a topic")
|
||||
parser.add_argument("--query-rewrite", action="store_true", help="Rewrite query before research for better results")
|
||||
parser.add_argument("--query-rewrite", action="store_true", help="Rewrite query for better results (works with any command)")
|
||||
parser.add_argument("--rewrite-tools", type=str, help="Tools for query rewriter (e.g., 'internet_search' or path to tools.py)")
|
||||
parser.add_argument("--tools", "-t", type=str, help="Path to tools.py file for research agent")
|
||||
parser.add_argument("--save", "-s", action="store_true", help="Save research output to file (output/research/)")
|
||||
@@ -748,10 +747,105 @@ class PraisonAI:
|
||||
|
||||
return args
|
||||
|
||||
def _rewrite_query(self, query: str, rewrite_tools: str = None, verbose: bool = False) -> str:
|
||||
"""
|
||||
Rewrite query using QueryRewriterAgent.
|
||||
|
||||
Args:
|
||||
query: The query to rewrite
|
||||
rewrite_tools: Tool names (comma-separated) or path to tools.py
|
||||
verbose: Enable verbose output
|
||||
|
||||
Returns:
|
||||
Rewritten query or original if rewriting fails
|
||||
"""
|
||||
try:
|
||||
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
|
||||
from rich import print
|
||||
|
||||
print("[bold cyan]Rewriting query for better results...[/bold cyan]")
|
||||
|
||||
# Load rewrite tools if specified
|
||||
rewrite_tools_list = []
|
||||
if rewrite_tools:
|
||||
if os.path.isfile(rewrite_tools):
|
||||
# Load from file
|
||||
try:
|
||||
import inspect
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location("rewrite_tools_module", rewrite_tools)
|
||||
if spec and spec.loader:
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
for name, obj in inspect.getmembers(module):
|
||||
if inspect.isfunction(obj) and not name.startswith('_'):
|
||||
rewrite_tools_list.append(obj)
|
||||
if rewrite_tools_list:
|
||||
print(f"[cyan]Loaded {len(rewrite_tools_list)} tools for query rewriter[/cyan]")
|
||||
except Exception as e:
|
||||
print(f"[yellow]Warning: Failed to load rewrite tools: {e}[/yellow]")
|
||||
else:
|
||||
# Treat as comma-separated tool names
|
||||
try:
|
||||
from praisonaiagents.tools import TOOL_MAPPINGS
|
||||
import praisonaiagents.tools as tools_module
|
||||
|
||||
tool_names = [t.strip() for t in rewrite_tools.split(',')]
|
||||
for tool_name in tool_names:
|
||||
if tool_name in TOOL_MAPPINGS:
|
||||
try:
|
||||
tool = getattr(tools_module, tool_name)
|
||||
rewrite_tools_list.append(tool)
|
||||
except Exception as e:
|
||||
print(f"[yellow]Warning: Failed to load rewrite tool '{tool_name}': {e}[/yellow]")
|
||||
else:
|
||||
print(f"[yellow]Warning: Unknown rewrite tool '{tool_name}'[/yellow]")
|
||||
if rewrite_tools_list:
|
||||
print(f"[cyan]Using rewrite tools: {', '.join(tool_names)}[/cyan]")
|
||||
except ImportError:
|
||||
print("[yellow]Warning: Could not import tools module[/yellow]")
|
||||
|
||||
rewriter = QueryRewriterAgent(
|
||||
model="gpt-4o-mini",
|
||||
verbose=verbose,
|
||||
tools=rewrite_tools_list if rewrite_tools_list else None
|
||||
)
|
||||
result = rewriter.rewrite(query, strategy=RewriteStrategy.AUTO)
|
||||
rewritten = result.primary_query
|
||||
|
||||
print(f"[cyan]Original:[/cyan] {query}")
|
||||
print(f"[cyan]Rewritten:[/cyan] {rewritten}")
|
||||
|
||||
return rewritten
|
||||
|
||||
except ImportError:
|
||||
from rich import print
|
||||
print("[yellow]Warning: QueryRewriterAgent not available, using original query[/yellow]")
|
||||
return query
|
||||
except Exception as e:
|
||||
from rich import print
|
||||
print(f"[yellow]Warning: Query rewrite failed ({e}), using original query[/yellow]")
|
||||
return query
|
||||
|
||||
def _rewrite_query_if_enabled(self, query: str) -> str:
|
||||
"""
|
||||
Rewrite query using QueryRewriterAgent if --query-rewrite is enabled.
|
||||
Returns the rewritten query or original if rewriting is disabled/fails.
|
||||
"""
|
||||
if not hasattr(self, 'args') or not getattr(self.args, 'query_rewrite', False):
|
||||
return query
|
||||
|
||||
rewrite_tools = getattr(self.args, 'rewrite_tools', None)
|
||||
verbose = getattr(self.args, 'verbose', False)
|
||||
return self._rewrite_query(query, rewrite_tools, verbose)
|
||||
|
||||
def handle_direct_prompt(self, prompt):
|
||||
"""
|
||||
Handle direct prompt by creating a single agent and running it.
|
||||
"""
|
||||
# Apply query rewriting if enabled
|
||||
prompt = self._rewrite_query_if_enabled(prompt)
|
||||
|
||||
if PRAISONAI_AVAILABLE:
|
||||
agent_config = {
|
||||
"name": "DirectAgent",
|
||||
@@ -994,60 +1088,8 @@ class PraisonAI:
|
||||
logging.getLogger('httpcore').setLevel(logging.WARNING)
|
||||
|
||||
# Rewrite query if requested
|
||||
original_query = query
|
||||
if query_rewrite:
|
||||
try:
|
||||
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
|
||||
print("[bold cyan]Rewriting query for better research results...[/bold cyan]")
|
||||
|
||||
# Load rewrite tools if specified
|
||||
rewrite_tools_list = []
|
||||
if rewrite_tools:
|
||||
if os.path.isfile(rewrite_tools):
|
||||
# Load from file
|
||||
try:
|
||||
import inspect
|
||||
spec = importlib.util.spec_from_file_location("rewrite_tools_module", rewrite_tools)
|
||||
if spec and spec.loader:
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
for name, obj in inspect.getmembers(module):
|
||||
if inspect.isfunction(obj) and not name.startswith('_'):
|
||||
rewrite_tools_list.append(obj)
|
||||
if rewrite_tools_list:
|
||||
print(f"[cyan]Loaded {len(rewrite_tools_list)} tools for query rewriter[/cyan]")
|
||||
except Exception as e:
|
||||
print(f"[yellow]Warning: Failed to load rewrite tools: {e}[/yellow]")
|
||||
else:
|
||||
# Treat as comma-separated tool names
|
||||
try:
|
||||
from praisonaiagents.tools import TOOL_MAPPINGS
|
||||
import praisonaiagents.tools as tools_module
|
||||
|
||||
tool_names = [t.strip() for t in rewrite_tools.split(',')]
|
||||
for tool_name in tool_names:
|
||||
if tool_name in TOOL_MAPPINGS:
|
||||
try:
|
||||
tool = getattr(tools_module, tool_name)
|
||||
rewrite_tools_list.append(tool)
|
||||
except Exception as e:
|
||||
print(f"[yellow]Warning: Failed to load rewrite tool '{tool_name}': {e}[/yellow]")
|
||||
else:
|
||||
print(f"[yellow]Warning: Unknown rewrite tool '{tool_name}'[/yellow]")
|
||||
if rewrite_tools_list:
|
||||
print(f"[cyan]Using rewrite tools: {', '.join(tool_names)}[/cyan]")
|
||||
except ImportError:
|
||||
print("[yellow]Warning: Could not import tools module[/yellow]")
|
||||
|
||||
rewriter = QueryRewriterAgent(model="gpt-4o-mini", verbose=verbose, tools=rewrite_tools_list if rewrite_tools_list else None)
|
||||
result = rewriter.rewrite(query, strategy=RewriteStrategy.AUTO)
|
||||
query = result.primary_query
|
||||
print(f"[cyan]Original:[/cyan] {original_query}")
|
||||
print(f"[cyan]Rewritten:[/cyan] {query}")
|
||||
except ImportError:
|
||||
print("[yellow]Warning: QueryRewriterAgent not available, using original query[/yellow]")
|
||||
except Exception as e:
|
||||
print(f"[yellow]Warning: Query rewrite failed ({e}), using original query[/yellow]")
|
||||
query = self._rewrite_query(query, rewrite_tools, verbose)
|
||||
|
||||
print("[bold green]Starting Deep Research...[/bold green]")
|
||||
print(f"Query: {query}")
|
||||
|
||||
@@ -57,7 +57,7 @@ class CloudDeployer:
|
||||
file.write("FROM python:3.11-slim\n")
|
||||
file.write("WORKDIR /app\n")
|
||||
file.write("COPY . .\n")
|
||||
file.write("RUN pip install flask praisonai==2.2.96 gunicorn markdown\n")
|
||||
file.write("RUN pip install flask praisonai==2.2.97 gunicorn markdown\n")
|
||||
file.write("EXPOSE 8080\n")
|
||||
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "2.2.96"
|
||||
__version__ = "2.2.97"
|
||||
@@ -12,7 +12,7 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=0.0.167",
|
||||
"praisonaiagents>=0.0.168",
|
||||
"python-dotenv>=0.19.0",
|
||||
"instructor>=1.3.3",
|
||||
"PyYAML>=6.0",
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
"""
|
||||
Tests for QueryRewriterAgent CLI integration
|
||||
|
||||
Run with: python -m pytest tests/unit/test_query_rewriter_cli.py -v
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import argparse
|
||||
|
||||
|
||||
class TestQueryRewriterCLIArgs:
|
||||
"""Test CLI argument parsing for query rewriter."""
|
||||
|
||||
def test_query_rewrite_arg_exists(self):
|
||||
"""Test --query-rewrite argument is available."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
# Create instance and check args can be parsed
|
||||
with patch('sys.argv', ['praisonai', '--help']):
|
||||
try:
|
||||
PraisonAI()
|
||||
except SystemExit:
|
||||
pass # --help causes exit, that's fine
|
||||
|
||||
def test_rewrite_tools_arg_exists(self):
|
||||
"""Test --rewrite-tools argument is available."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
with patch('sys.argv', ['praisonai', '--help']):
|
||||
try:
|
||||
PraisonAI()
|
||||
except SystemExit:
|
||||
pass
|
||||
|
||||
|
||||
class TestRewriteQueryMethod:
|
||||
"""Test the _rewrite_query helper method."""
|
||||
|
||||
def test_rewrite_query_returns_rewritten(self):
|
||||
"""Test _rewrite_query returns rewritten query."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
# Setup mock
|
||||
mock_result = MagicMock()
|
||||
mock_result.primary_query = "What are the current trends in artificial intelligence?"
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.rewrite.return_value = mock_result
|
||||
|
||||
# Patch at the import location inside the method
|
||||
with patch.dict('sys.modules', {'praisonaiagents': MagicMock()}):
|
||||
import sys
|
||||
sys.modules['praisonaiagents'].QueryRewriterAgent = MagicMock(return_value=mock_agent)
|
||||
sys.modules['praisonaiagents'].RewriteStrategy = MagicMock()
|
||||
|
||||
praison = PraisonAI()
|
||||
result = praison._rewrite_query("AI trends", None, False)
|
||||
|
||||
assert result == "What are the current trends in artificial intelligence?"
|
||||
|
||||
def test_rewrite_query_fallback_on_import_error(self):
|
||||
"""Test _rewrite_query returns original on ImportError."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
# Mock the import to raise ImportError
|
||||
praison = PraisonAI()
|
||||
|
||||
with patch.object(praison, '_rewrite_query') as mock_method:
|
||||
# Simulate the actual behavior - returns original on error
|
||||
mock_method.return_value = "AI trends"
|
||||
result = mock_method("AI trends", None, False)
|
||||
assert result == "AI trends"
|
||||
|
||||
def test_rewrite_query_fallback_on_exception(self):
|
||||
"""Test _rewrite_query returns original on exception."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
# Patch to raise exception
|
||||
with patch.dict('sys.modules', {'praisonaiagents': MagicMock()}):
|
||||
import sys
|
||||
sys.modules['praisonaiagents'].QueryRewriterAgent = MagicMock(side_effect=Exception("Test error"))
|
||||
|
||||
praison = PraisonAI()
|
||||
result = praison._rewrite_query("AI trends", None, False)
|
||||
|
||||
# Should return original query on error
|
||||
assert result == "AI trends"
|
||||
|
||||
|
||||
class TestRewriteQueryIfEnabled:
|
||||
"""Test the _rewrite_query_if_enabled wrapper method."""
|
||||
|
||||
def test_returns_original_when_disabled(self):
|
||||
"""Test returns original query when --query-rewrite not set."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
praison = PraisonAI()
|
||||
praison.args = argparse.Namespace(query_rewrite=False)
|
||||
|
||||
result = praison._rewrite_query_if_enabled("test query")
|
||||
assert result == "test query"
|
||||
|
||||
def test_returns_original_when_no_args(self):
|
||||
"""Test returns original query when args not set."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
praison = PraisonAI()
|
||||
# Don't set args at all
|
||||
if hasattr(praison, 'args'):
|
||||
delattr(praison, 'args')
|
||||
|
||||
result = praison._rewrite_query_if_enabled("test query")
|
||||
assert result == "test query"
|
||||
|
||||
def test_calls_rewrite_when_enabled(self):
|
||||
"""Test calls _rewrite_query when --query-rewrite is set."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
praison = PraisonAI()
|
||||
praison.args = argparse.Namespace(
|
||||
query_rewrite=True,
|
||||
rewrite_tools=None,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
with patch.object(praison, '_rewrite_query', return_value="rewritten query") as mock_rewrite:
|
||||
result = praison._rewrite_query_if_enabled("original query")
|
||||
|
||||
mock_rewrite.assert_called_once_with("original query", None, False)
|
||||
assert result == "rewritten query"
|
||||
|
||||
def test_passes_tools_and_verbose(self):
|
||||
"""Test passes rewrite_tools and verbose to _rewrite_query."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
praison = PraisonAI()
|
||||
praison.args = argparse.Namespace(
|
||||
query_rewrite=True,
|
||||
rewrite_tools="internet_search",
|
||||
verbose=True
|
||||
)
|
||||
|
||||
with patch.object(praison, '_rewrite_query', return_value="rewritten") as mock_rewrite:
|
||||
praison._rewrite_query_if_enabled("query")
|
||||
|
||||
mock_rewrite.assert_called_once_with("query", "internet_search", True)
|
||||
|
||||
|
||||
class TestHandleDirectPromptWithRewrite:
|
||||
"""Test handle_direct_prompt integrates query rewriting."""
|
||||
|
||||
def test_applies_rewrite_before_agent(self):
|
||||
"""Test query is rewritten before passing to agent."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
praison = PraisonAI()
|
||||
praison.args = argparse.Namespace(
|
||||
query_rewrite=True,
|
||||
rewrite_tools=None,
|
||||
verbose=False,
|
||||
llm=None
|
||||
)
|
||||
|
||||
with patch.object(praison, '_rewrite_query_if_enabled', return_value="rewritten prompt") as mock_rewrite:
|
||||
with patch('praisonai.cli.PRAISONAI_AVAILABLE', True):
|
||||
with patch('praisonai.cli.PraisonAgent') as mock_agent_class:
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.start.return_value = "result"
|
||||
mock_agent_class.return_value = mock_agent
|
||||
|
||||
praison.handle_direct_prompt("original prompt")
|
||||
|
||||
# Verify rewrite was called
|
||||
mock_rewrite.assert_called_once_with("original prompt")
|
||||
|
||||
# Verify agent received rewritten prompt
|
||||
mock_agent.start.assert_called_once_with("rewritten prompt")
|
||||
|
||||
|
||||
class TestToolLoading:
|
||||
"""Test tool loading for query rewriter."""
|
||||
|
||||
def test_load_builtin_tool_by_name(self):
|
||||
"""Test loading built-in tool by name."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
# Setup mocks
|
||||
mock_result = MagicMock()
|
||||
mock_result.primary_query = "rewritten"
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.rewrite.return_value = mock_result
|
||||
|
||||
mock_tool = MagicMock()
|
||||
mock_tools_module = MagicMock()
|
||||
mock_tools_module.TOOL_MAPPINGS = {'internet_search': 'praisonaiagents.tools'}
|
||||
mock_tools_module.internet_search = mock_tool
|
||||
|
||||
mock_praisonaiagents = MagicMock()
|
||||
mock_praisonaiagents.QueryRewriterAgent = MagicMock(return_value=mock_agent)
|
||||
mock_praisonaiagents.RewriteStrategy = MagicMock()
|
||||
mock_praisonaiagents.tools = mock_tools_module
|
||||
|
||||
with patch.dict('sys.modules', {
|
||||
'praisonaiagents': mock_praisonaiagents,
|
||||
'praisonaiagents.tools': mock_tools_module
|
||||
}):
|
||||
praison = PraisonAI()
|
||||
result = praison._rewrite_query("query", "internet_search", False)
|
||||
|
||||
# Verify result
|
||||
assert result == "rewritten"
|
||||
|
||||
def test_handles_unknown_tool_gracefully(self):
|
||||
"""Test handles unknown tool name gracefully."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
# Setup mocks
|
||||
mock_result = MagicMock()
|
||||
mock_result.primary_query = "rewritten"
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.rewrite.return_value = mock_result
|
||||
|
||||
mock_tools_module = MagicMock()
|
||||
mock_tools_module.TOOL_MAPPINGS = {} # Empty - no tools
|
||||
|
||||
mock_praisonaiagents = MagicMock()
|
||||
mock_praisonaiagents.QueryRewriterAgent = MagicMock(return_value=mock_agent)
|
||||
mock_praisonaiagents.RewriteStrategy = MagicMock()
|
||||
mock_praisonaiagents.tools = mock_tools_module
|
||||
|
||||
with patch.dict('sys.modules', {
|
||||
'praisonaiagents': mock_praisonaiagents,
|
||||
'praisonaiagents.tools': mock_tools_module
|
||||
}):
|
||||
praison = PraisonAI()
|
||||
# Should not raise, just warn
|
||||
result = praison._rewrite_query("query", "unknown_tool", False)
|
||||
|
||||
assert result == "rewritten"
|
||||
|
||||
|
||||
class TestQueryRewriterIntegration:
|
||||
"""Integration tests for QueryRewriterAgent with CLI."""
|
||||
|
||||
def test_full_rewrite_flow_mocked(self):
|
||||
"""Test full rewrite flow with mocked components."""
|
||||
from praisonai.cli import PraisonAI
|
||||
|
||||
# Setup mocks
|
||||
mock_result = MagicMock()
|
||||
mock_result.primary_query = "What are the latest developments in AI technology?"
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.rewrite.return_value = mock_result
|
||||
|
||||
mock_praisonaiagents = MagicMock()
|
||||
mock_praisonaiagents.QueryRewriterAgent = MagicMock(return_value=mock_agent)
|
||||
mock_praisonaiagents.RewriteStrategy = MagicMock()
|
||||
|
||||
with patch.dict('sys.modules', {'praisonaiagents': mock_praisonaiagents}):
|
||||
praison = PraisonAI()
|
||||
praison.args = argparse.Namespace(
|
||||
query_rewrite=True,
|
||||
rewrite_tools=None,
|
||||
verbose=False,
|
||||
llm=None
|
||||
)
|
||||
|
||||
result = praison._rewrite_query_if_enabled("AI trends")
|
||||
|
||||
assert "AI" in result or "developments" in result
|
||||
mock_agent.rewrite.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Generated
+4
-4
@@ -4197,7 +4197,7 @@ requires-dist = [
|
||||
{ name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.0.22" },
|
||||
{ name = "praisonai-tools", marker = "extra == 'autogen-v4'", specifier = ">=0.0.22" },
|
||||
{ name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.0.22" },
|
||||
{ name = "praisonaiagents", specifier = ">=0.0.167" },
|
||||
{ name = "praisonaiagents", specifier = ">=0.0.168" },
|
||||
{ name = "pyautogen", marker = "extra == 'autogen'", specifier = "==0.2.29" },
|
||||
{ name = "pydantic", marker = "extra == 'chat'", specifier = "<=2.10.1" },
|
||||
{ name = "pydantic", marker = "extra == 'code'", specifier = "<=2.10.1" },
|
||||
@@ -4252,7 +4252,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.0.167"
|
||||
version = "0.0.168"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -4261,9 +4261,9 @@ dependencies = [
|
||||
{ name = "pydantic" },
|
||||
{ name = "rich" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/8e/26/2b77721ee85f5943ea5f1515453a32bfb7f6442128f5eb2ce964d2016d46/praisonaiagents-0.0.167.tar.gz", hash = "sha256:fd052746bcb40f14886f58786ec06cc6efa7f564e1ff0dea7802bface72b560b", size = 304811, upload-time = "2025-12-15T07:56:09.8Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/71/a9/d4e0effb144fefe14d724dea9765cf27935d6bd9cd1c76b80ac8961d2f86/praisonaiagents-0.0.168.tar.gz", hash = "sha256:744156db568d595004d8b29a72a214dbad6a9a4702b6f72eff8f08aba77515b1", size = 304818, upload-time = "2025-12-15T08:10:54.72Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1b/3c/5297203ed5385b2fa2c9e7b3576a7ff685a72fd4c431bad6d2768ea5f03c/praisonaiagents-0.0.167-py3-none-any.whl", hash = "sha256:fee9645ed990b060a20aa95caf64d4524b0e31c852c46b2eb8344351886210aa", size = 322038, upload-time = "2025-12-15T07:56:08.184Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7b/42/3077b1892c3ec0cce059a17dd8517e8f7b4b7944115e3d6025eeb91f633c/praisonaiagents-0.0.168-py3-none-any.whl", hash = "sha256:3e7095da3ebbe1d7c23613418b3200f3d42198ae2ccde08dd753cf548a549e93", size = 322039, upload-time = "2025-12-15T08:10:53.526Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in new issue
Block a user