Compare commits

...
5 Commits
25 changed files with 2063 additions and 230 deletions

No files matched your search

+65
View File
@@ -54,6 +54,7 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
- 📄 YAML Configuration
- 💯 100+ LLM Support
- 🔬 **Deep Research Agents** (OpenAI & Gemini)
- 🔄 **Query Rewriter Agent** (HyDE, Step-back, Multi-query)
## Using Python Code
@@ -135,6 +136,41 @@ print(result.report)
- 🛠️ Built-in tools: web search, code interpreter, MCP, file search
- 🔄 Automatic provider detection from model name
### 4. Query Rewriter Agent
Transform user queries to improve RAG retrieval quality using multiple strategies.
```python
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent(model="gpt-4o-mini")
# Basic - expands abbreviations, adds context
result = agent.rewrite("AI trends")
print(result.primary_query) # "What are the current trends in Artificial Intelligence?"
# HyDE - generates hypothetical document for semantic matching
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
# Step-back - generates broader context question
result = agent.rewrite("GPT-4 vs Claude 3?", strategy=RewriteStrategy.STEP_BACK)
# Sub-queries - decomposes complex questions
result = agent.rewrite("RAG setup and best embedding models?", strategy=RewriteStrategy.SUB_QUERIES)
# Contextual - resolves references using chat history
result = agent.rewrite("What about cost?", chat_history=[...])
```
**Strategies:**
- **BASIC**: Expand abbreviations, fix typos, add context
- **HYDE**: Generate hypothetical document for semantic matching
- **STEP_BACK**: Generate higher-level concept questions
- **SUB_QUERIES**: Decompose multi-part questions
- **MULTI_QUERY**: Generate multiple paraphrased versions
- **CONTEXTUAL**: Resolve references using conversation history
- **AUTO**: Automatically detect best strategy
## Using No Code
### Auto Mode:
@@ -144,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)
@@ -152,10 +200,27 @@ praisonai research "What are the latest AI trends in 2025?"
# Use Gemini
praisonai research --model deep-research-pro "Your research query"
# Rewrite query before research
praisonai research --query-rewrite "AI trends"
# Rewrite with search tools
praisonai research --query-rewrite --rewrite-tools "internet_search" "AI trends"
# Use custom tools from file (gathers context before deep research)
praisonai research --tools tools.py "Your research query"
praisonai research -t my_tools.py "Your research query"
# Use built-in tools by name (comma-separated)
praisonai research --tools "internet_search,wiki_search" "Your query"
praisonai research -t "yfinance,calculator_tools" "Stock analysis query"
# Save output to file (output/research/{query}.md)
praisonai research --save "Your research query"
praisonai research -s "Your research query"
# Combine options
praisonai research --query-rewrite --tools tools.py --save "Your research query"
# Verbose mode (show debug logs)
praisonai research -v "Your research query"
```
+1 -1
View File
@@ -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.93" \
"praisonai>=2.2.97" \
"praisonai[api]" \
gunicorn \
markdown
+1 -1
View File
@@ -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.93" \
"praisonai>=2.2.97" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -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.93" \
"praisonai>=2.2.97" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -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.93" \
"praisonai>=2.2.97" \
"praisonai[ui]" \
"praisonai[crewai]"
+2 -2
View File
@@ -121,7 +121,7 @@ healthcheck:
## 📦 Package Versions
All Docker images use consistent, up-to-date versions:
- PraisonAI: `>=2.2.93`
- PraisonAI: `>=2.2.95`
- PraisonAI Agents: `>=0.0.92`
- Python: `3.11-slim`
@@ -218,7 +218,7 @@ docker-compose up -d
### Version Pinning
To use specific versions, update the Dockerfile:
```dockerfile
RUN pip install "praisonai==2.2.93" "praisonaiagents==0.0.92"
RUN pip install "praisonai==2.2.95" "praisonaiagents==0.0.92"
```
## 🌐 Production Deployment
+41 -3
View File
@@ -136,7 +136,45 @@ result = agent.research("Research quantum computing advances")
- **Citations**: Structured citations with URLs
- **Tools**: Web search, code interpreter, MCP, file search
### 3. ImageAgent
### 3. QueryRewriterAgent 🆕
Transform user queries to improve RAG retrieval quality using multiple strategies.
```python
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent(model="gpt-4o-mini", verbose=True)
# Basic rewriting - expands abbreviations, adds context
result = agent.rewrite("AI trends")
print(result.primary_query)
# Output: "What are the current trends in Artificial Intelligence (AI)?"
# HyDE - generates hypothetical document for semantic matching
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
print(result.hypothetical_document)
# Step-back - generates broader context question
result = agent.rewrite("Difference between GPT-4 and Claude 3?", strategy=RewriteStrategy.STEP_BACK)
print(result.step_back_question)
# Sub-queries - decomposes complex questions
result = agent.rewrite("How to set up RAG and what embedding models to use?", strategy=RewriteStrategy.SUB_QUERIES)
print(result.sub_queries)
# Contextual - uses chat history to resolve references
result = agent.rewrite("What about its cost?", strategy=RewriteStrategy.CONTEXTUAL, chat_history=[...])
```
**Strategies:**
- **BASIC**: Expand abbreviations, fix typos, add context
- **HYDE**: Generate hypothetical document for better semantic matching
- **STEP_BACK**: Generate higher-level concept questions
- **SUB_QUERIES**: Decompose multi-part questions
- **MULTI_QUERY**: Generate multiple paraphrased versions
- **CONTEXTUAL**: Resolve references using conversation history
- **AUTO**: Automatically detect best strategy
### 4. ImageAgent
Agent specialized for image generation and analysis.
```python
@@ -150,7 +188,7 @@ agent = ImageAgent(
result = agent.chat("Create an image of a sunset over mountains")
```
### 4. ContextAgent
### 5. ContextAgent
Agent with persistent context across conversations.
```python
@@ -162,7 +200,7 @@ agent = ContextAgent(
)
```
### 5. RouterAgent
### 6. RouterAgent
Routes requests to appropriate specialized agents.
```python
@@ -26,6 +26,11 @@ from .agent.deep_research_agent import (
FileSearchCall,
Provider
)
from .agent.query_rewriter_agent import (
QueryRewriterAgent,
RewriteStrategy,
RewriteResult
)
from .agents.agents import PraisonAIAgents
from .task.task import Task
from .tools.tools import Tools
@@ -206,7 +211,10 @@ __all__ = [
'CodeExecutionStep',
'MCPCall',
'FileSearchCall',
'Provider'
'Provider',
'QueryRewriterAgent',
'RewriteStrategy',
'RewriteResult'
]
# Add MCP to __all__ if available
@@ -15,6 +15,11 @@ from .deep_research_agent import (
FileSearchCall,
Provider
)
from .query_rewriter_agent import (
QueryRewriterAgent,
RewriteStrategy,
RewriteResult
)
__all__ = [
'Agent',
@@ -35,5 +40,8 @@ __all__ = [
'CodeExecutionStep',
'MCPCall',
'FileSearchCall',
'Provider'
'Provider',
'QueryRewriterAgent',
'RewriteStrategy',
'RewriteResult'
]
@@ -0,0 +1,627 @@
"""
Query Rewriter Agent Module
This module provides the QueryRewriterAgent class for transforming user queries
to improve retrieval quality in RAG applications.
Supported Rewriting Strategies:
- **BASIC**: Simple rephrasing for clarity and keyword optimization
- **HYDE**: Hypothetical Document Embeddings - generates a hypothetical answer
- **STEP_BACK**: Generates higher-level concept questions for complex queries
- **SUB_QUERIES**: Decomposes multi-part questions into focused sub-queries
- **MULTI_QUERY**: Generates multiple paraphrased versions for ensemble retrieval
- **CONTEXTUAL**: Uses conversation history to resolve references and context
Example:
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent(verbose=True)
# Basic rewriting
result = agent.rewrite("AI trends")
print(result.rewritten_queries)
# HyDE for better semantic matching
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
print(result.hypothetical_document)
# Contextual with chat history
chat_history = [
{"role": "user", "content": "Tell me about Python"},
{"role": "assistant", "content": "Python is a programming language..."}
]
result = agent.rewrite("What about its performance?", chat_history=chat_history)
# With search tools for context-aware rewriting
from praisonaiagents.tools import internet_search
agent = QueryRewriterAgent(tools=[internet_search], verbose=True)
result = agent.rewrite("latest AI developments") # Searches first, then rewrites
"""
import logging
import json
from typing import List, Optional, Any, Dict
from dataclasses import dataclass, field
from enum import Enum
class RewriteStrategy(Enum):
"""Enumeration of available query rewriting strategies."""
BASIC = "basic"
HYDE = "hyde"
STEP_BACK = "step_back"
SUB_QUERIES = "sub_queries"
MULTI_QUERY = "multi_query"
CONTEXTUAL = "contextual"
AUTO = "auto"
@dataclass
class RewriteResult:
"""Result of a query rewriting operation."""
original_query: str
rewritten_queries: List[str]
strategy_used: RewriteStrategy
hypothetical_document: Optional[str] = None
step_back_question: Optional[str] = None
sub_queries: Optional[List[str]] = None
metadata: Dict[str, Any] = field(default_factory=dict)
def __repr__(self):
return f"RewriteResult(strategy={self.strategy_used.value}, queries={len(self.rewritten_queries)})"
@property
def primary_query(self) -> str:
"""Returns the primary rewritten query."""
return self.rewritten_queries[0] if self.rewritten_queries else self.original_query
@property
def all_queries(self) -> List[str]:
"""Returns all queries including original and rewritten."""
queries = [self.original_query]
queries.extend(self.rewritten_queries)
if self.step_back_question:
queries.append(self.step_back_question)
if self.sub_queries:
queries.extend(self.sub_queries)
return list(set(queries)) # Remove duplicates
class QueryRewriterAgent:
"""
Agent for rewriting queries to improve retrieval quality in RAG applications.
This agent transforms user queries using various strategies to bridge the gap
between how users ask questions and how information is stored in knowledge bases.
Attributes:
name: Name of the agent
model: LLM model to use for rewriting
verbose: Whether to print detailed logs
max_queries: Maximum number of queries to generate for multi-query strategy
abbreviations: Dictionary of common abbreviations to expand
Example:
agent = QueryRewriterAgent(model="gpt-4o-mini", verbose=True)
result = agent.rewrite("ML best practices", strategy=RewriteStrategy.MULTI_QUERY)
for query in result.rewritten_queries:
print(query)
"""
# Default prompts for each strategy
PROMPTS = {
"basic": """You are an expert at rewriting search queries for better retrieval.
Given the following query, rewrite it to be clearer, more specific, and better suited for searching a knowledge base.
Guidelines:
- Expand abbreviations and acronyms
- Fix typos and grammatical errors
- Add relevant context if the query is too short
- Transform keyword searches into natural questions
- Keep the original intent intact
Query: {query}
Rewritten query:""",
"hyde": """You are an expert at generating hypothetical documents.
Given the following query, write a detailed hypothetical document that would perfectly answer this query. This document will be used to find similar real documents in a knowledge base.
Guidelines:
- Write as if you are answering the query directly
- Include specific details, facts, and terminology
- Use the same style and vocabulary as the target documents
- Be comprehensive but focused on the query topic
Query: {query}
Hypothetical document:""",
"step_back": """You are an expert at generating step-back questions.
Given the following query, generate a more general, higher-level question that captures the broader concept. This helps retrieve background information needed to answer the specific query.
Guidelines:
- Identify the underlying concept or principle
- Create a question about fundamentals or general knowledge
- The step-back question should provide context for the original query
Query: {query}
Step-back question:""",
"sub_queries": """You are an expert at decomposing complex queries.
Given the following query, break it down into simpler, focused sub-queries. Each sub-query should address one specific aspect of the original question.
Guidelines:
- Identify distinct aspects or components of the query
- Create independent, self-contained sub-queries
- Each sub-query should be answerable on its own
- Return as a JSON array of strings
Query: {query}
Sub-queries (JSON array):""",
"multi_query": """You are an expert at query expansion and paraphrasing.
Given the following query, generate {num_queries} different versions that capture the same intent but use different words, phrases, or perspectives.
Guidelines:
- Use synonyms and alternative phrasings
- Try different question formats (what, how, why, etc.)
- Include both formal and informal versions
- Add relevant context or specificity
- Return as a JSON array of strings
Query: {query}
Alternative queries (JSON array):""",
"contextual": """You are an expert at understanding conversational context.
Given the conversation history and a follow-up query, rewrite the query as a standalone question that captures all relevant context.
Guidelines:
- Resolve pronouns and references (it, this, that, etc.)
- Include relevant context from the conversation
- Keep entity names and important details
- Make the query self-contained and clear
Conversation history:
{chat_history}
Follow-up query: {query}
Standalone question:""",
"auto": """You are an expert at analyzing and rewriting queries.
Analyze the following query and determine the best rewriting strategy, then apply it.
Query types to consider:
1. Short/keyword queries → Expand into full questions
2. Ambiguous queries → Clarify with context
3. Complex multi-part queries → Break into sub-queries
4. Follow-up queries → Resolve references from context
5. Well-formed queries → Minor optimization only
Query: {query}
{context}
Analysis and rewritten query:"""
}
def __init__(
self,
name: str = "QueryRewriterAgent",
model: str = "gpt-4o-mini",
instructions: Optional[str] = None,
verbose: bool = False,
max_queries: int = 5,
abbreviations: Optional[Dict[str, str]] = None,
temperature: float = 0.3,
max_tokens: int = 500,
tools: Optional[List[Any]] = None
):
"""
Initialize the QueryRewriterAgent.
Args:
name: Name of the agent
model: LLM model to use (default: gpt-4o-mini)
instructions: Custom instructions for the agent
verbose: Whether to print detailed logs
max_queries: Maximum queries for multi-query strategy
abbreviations: Custom abbreviation expansions
temperature: Temperature for LLM generation
max_tokens: Maximum tokens for LLM response
tools: Optional list of tools (agent decides when to use them)
"""
self.name = name
self.model = model
self.instructions = instructions
self.verbose = verbose
self.max_queries = max_queries
self.temperature = temperature
self.max_tokens = max_tokens
self.tools = tools or []
self._agent = None # Lazy initialized
# Default abbreviations (can be extended)
self.abbreviations = abbreviations or {
"AI": "Artificial Intelligence",
"ML": "Machine Learning",
"NLP": "Natural Language Processing",
"LLM": "Large Language Model",
"RAG": "Retrieval Augmented Generation",
"API": "Application Programming Interface",
"DB": "Database",
"UI": "User Interface",
"UX": "User Experience",
"SDK": "Software Development Kit",
"CLI": "Command Line Interface",
}
# Setup logging
self.logger = logging.getLogger(__name__)
if self.verbose:
self.logger.setLevel(logging.DEBUG)
@property
def agent(self):
"""Lazy initialization of internal Agent."""
if self._agent is None:
from .agent import Agent
self._agent = Agent(
name=self.name,
role="Query Rewriting Specialist",
goal="Rewrite queries for better search and retrieval results. Use available tools when needed to gather context.",
backstory="You are an expert at understanding user intent and transforming queries for optimal information retrieval.",
tools=self.tools,
llm=self.model,
verbose=self.verbose,
markdown=False
)
return self._agent
def _call_agent(self, prompt: str) -> str:
"""Call the internal Agent with the given prompt."""
try:
response = self.agent.chat(prompt)
return response if response else ""
except Exception as e:
self.logger.error(f"Agent call failed: {e}")
raise
def _expand_abbreviations(self, query: str) -> str:
"""Expand known abbreviations in the query."""
words = query.split()
expanded = []
for word in words:
clean_word = word.strip(".,!?;:")
if clean_word.upper() in self.abbreviations:
expanded.append(f"{clean_word} ({self.abbreviations[clean_word.upper()]})")
else:
expanded.append(word)
return " ".join(expanded)
def _is_short_query(self, query: str, threshold: int = 5) -> bool:
"""Check if query is considered short (keyword-style)."""
return len(query.split()) < threshold
def _parse_json_array(self, text: str) -> List[str]:
"""Parse a JSON array from LLM response."""
try:
# Try direct JSON parsing
result = json.loads(text)
if isinstance(result, list):
return [str(item) for item in result]
except json.JSONDecodeError:
pass
# Try to extract JSON array from text
import re
match = re.search(r'\[.*?\]', text, re.DOTALL)
if match:
try:
result = json.loads(match.group())
if isinstance(result, list):
return [str(item) for item in result]
except json.JSONDecodeError:
pass
# Fallback: split by newlines and clean up
lines = [line.strip() for line in text.split('\n') if line.strip()]
lines = [line.lstrip('0123456789.-) ').strip('"\'') for line in lines]
return [line for line in lines if line and len(line) > 3]
def rewrite(
self,
query: str,
strategy: RewriteStrategy = RewriteStrategy.AUTO,
chat_history: Optional[List[Dict[str, str]]] = None,
context: Optional[str] = None,
num_queries: int = None
) -> RewriteResult:
"""
Rewrite a query using the specified strategy.
All operations go through an internal Agent which automatically handles
tool calling when tools are provided. The agent decides when to use tools
based on the query context.
Args:
query: The original user query
strategy: Rewriting strategy to use
chat_history: Previous conversation messages for contextual rewriting
context: Additional context about the knowledge base
num_queries: Number of queries for multi-query strategy
Returns:
RewriteResult containing rewritten queries
"""
if self.verbose:
self.logger.debug(f"Rewriting query: '{query}' with strategy: {strategy.value}")
if self.tools:
print(f"[cyan]Agent has {len(self.tools)} tools available (will use if needed)...[/cyan]")
# Auto-detect strategy if needed
if strategy == RewriteStrategy.AUTO:
strategy = self._detect_strategy(query, chat_history)
if self.verbose:
self.logger.debug(f"Auto-detected strategy: {strategy.value}")
# All calls go through the internal Agent which handles tool calling
combined_context = context or ""
# Dispatch to appropriate method
if strategy == RewriteStrategy.BASIC:
return self._rewrite_basic(query, combined_context)
elif strategy == RewriteStrategy.HYDE:
return self._rewrite_hyde(query, combined_context)
elif strategy == RewriteStrategy.STEP_BACK:
return self._rewrite_step_back(query, combined_context)
elif strategy == RewriteStrategy.SUB_QUERIES:
return self._rewrite_sub_queries(query, combined_context)
elif strategy == RewriteStrategy.MULTI_QUERY:
return self._rewrite_multi_query(query, num_queries or self.max_queries, combined_context)
elif strategy == RewriteStrategy.CONTEXTUAL:
return self._rewrite_contextual(query, chat_history or [], combined_context)
else:
# Default to basic
return self._rewrite_basic(query, combined_context)
def _detect_strategy(
self,
query: str,
chat_history: Optional[List[Dict[str, str]]] = None
) -> RewriteStrategy:
"""Auto-detect the best rewriting strategy for a query."""
# If chat history exists and query seems like a follow-up
if chat_history and self._is_follow_up(query):
return RewriteStrategy.CONTEXTUAL
# Short keyword queries
if self._is_short_query(query):
return RewriteStrategy.BASIC
# Complex multi-part queries
if self._is_complex_query(query):
return RewriteStrategy.SUB_QUERIES
# Questions that might benefit from step-back
if self._needs_step_back(query):
return RewriteStrategy.STEP_BACK
# Default to basic
return RewriteStrategy.BASIC
def _is_follow_up(self, query: str) -> bool:
"""Check if query appears to be a follow-up question."""
follow_up_indicators = [
"it", "this", "that", "they", "them", "their",
"what about", "how about", "and", "also", "more",
"same", "similar", "related"
]
query_lower = query.lower()
return any(indicator in query_lower for indicator in follow_up_indicators)
def _is_complex_query(self, query: str) -> bool:
"""Check if query contains multiple distinct questions."""
# Check for multiple question marks or conjunctions
if query.count('?') > 1:
return True
complex_indicators = [" and ", " or ", " also ", " as well as ", " plus "]
return any(ind in query.lower() for ind in complex_indicators)
def _needs_step_back(self, query: str) -> bool:
"""Check if query would benefit from a step-back question."""
# Queries about specific details often benefit from broader context
specific_indicators = [
"specific", "exactly", "precisely", "particular",
"difference between", "compare", "versus", "vs"
]
return any(ind in query.lower() for ind in specific_indicators)
def _rewrite_basic(self, query: str, context: str = "") -> RewriteResult:
"""Basic query rewriting for clarity and keyword optimization."""
# First expand abbreviations
expanded_query = self._expand_abbreviations(query)
# Build prompt with optional context
prompt = self.PROMPTS["basic"].format(query=expanded_query)
if context:
prompt += f"\n\nAdditional context to help with rewriting:\n{context}"
rewritten = self._call_agent(prompt).strip()
return RewriteResult(
original_query=query,
rewritten_queries=[rewritten],
strategy_used=RewriteStrategy.BASIC,
metadata={"expanded_abbreviations": expanded_query != query, "has_context": bool(context)}
)
def _rewrite_hyde(self, query: str, context: str = "") -> RewriteResult:
"""HyDE: Generate hypothetical document for better semantic matching."""
prompt = self.PROMPTS["hyde"].format(query=query)
if context:
prompt += f"\n\nUse this context to make the hypothetical document more accurate:\n{context}"
hypothetical_doc = self._call_agent(prompt).strip()
# The hypothetical document itself becomes the search query
return RewriteResult(
original_query=query,
rewritten_queries=[hypothetical_doc[:500]], # Truncate for embedding
strategy_used=RewriteStrategy.HYDE,
hypothetical_document=hypothetical_doc,
metadata={"full_document_length": len(hypothetical_doc), "has_context": bool(context)}
)
def _rewrite_step_back(self, query: str, context: str = "") -> RewriteResult:
"""Generate step-back question for broader context retrieval."""
# Generate step-back question
prompt = self.PROMPTS["step_back"].format(query=query)
if context:
prompt += f"\n\nContext:\n{context}"
step_back = self._call_agent(prompt).strip()
# Also do basic rewriting of original
basic_result = self._rewrite_basic(query, context)
return RewriteResult(
original_query=query,
rewritten_queries=basic_result.rewritten_queries,
strategy_used=RewriteStrategy.STEP_BACK,
step_back_question=step_back,
metadata={"includes_step_back": True, "has_context": bool(context)}
)
def _rewrite_sub_queries(self, query: str, context: str = "") -> RewriteResult:
"""Decompose complex query into focused sub-queries."""
prompt = self.PROMPTS["sub_queries"].format(query=query)
if context:
prompt += f"\n\nContext to help with decomposition:\n{context}"
response = self._call_agent(prompt).strip()
sub_queries = self._parse_json_array(response)
# Ensure we have at least one sub-query
if not sub_queries:
sub_queries = [query]
return RewriteResult(
original_query=query,
rewritten_queries=sub_queries,
strategy_used=RewriteStrategy.SUB_QUERIES,
sub_queries=sub_queries,
metadata={"num_sub_queries": len(sub_queries), "has_context": bool(context)}
)
def _rewrite_multi_query(self, query: str, num_queries: int, context: str = "") -> RewriteResult:
"""Generate multiple paraphrased versions for ensemble retrieval."""
prompt = self.PROMPTS["multi_query"].format(
query=query,
num_queries=num_queries
)
if context:
prompt += f"\n\nContext to help with paraphrasing:\n{context}"
response = self._call_agent(prompt).strip()
queries = self._parse_json_array(response)
# Ensure we have queries and limit to max
if not queries:
queries = [query]
queries = queries[:num_queries]
return RewriteResult(
original_query=query,
rewritten_queries=queries,
strategy_used=RewriteStrategy.MULTI_QUERY,
metadata={"num_generated": len(queries), "has_context": bool(context)}
)
def _rewrite_contextual(
self,
query: str,
chat_history: List[Dict[str, str]],
context: str = ""
) -> RewriteResult:
"""Rewrite query using conversation context."""
# Format chat history
history_str = ""
for msg in chat_history[-10:]: # Last 10 messages
role = msg.get("role", "user")
content = msg.get("content", "")
history_str += f"{role.capitalize()}: {content}\n"
prompt = self.PROMPTS["contextual"].format(
chat_history=history_str,
query=query
)
if context:
prompt += f"\n\nAdditional context:\n{context}"
rewritten = self._call_agent(prompt).strip()
return RewriteResult(
original_query=query,
rewritten_queries=[rewritten],
strategy_used=RewriteStrategy.CONTEXTUAL,
metadata={
"history_length": len(chat_history),
"history_used": min(10, len(chat_history)),
"has_context": bool(context)
}
)
# Convenience methods for direct strategy access
def rewrite_basic(self, query: str) -> RewriteResult:
"""Convenience method for basic rewriting."""
return self.rewrite(query, strategy=RewriteStrategy.BASIC)
def rewrite_hyde(self, query: str) -> RewriteResult:
"""Convenience method for HyDE rewriting."""
return self.rewrite(query, strategy=RewriteStrategy.HYDE)
def rewrite_step_back(self, query: str) -> RewriteResult:
"""Convenience method for step-back rewriting."""
return self.rewrite(query, strategy=RewriteStrategy.STEP_BACK)
def rewrite_sub_queries(self, query: str) -> RewriteResult:
"""Convenience method for sub-query decomposition."""
return self.rewrite(query, strategy=RewriteStrategy.SUB_QUERIES)
def rewrite_multi_query(self, query: str, num_queries: int = None) -> RewriteResult:
"""Convenience method for multi-query generation."""
return self.rewrite(
query,
strategy=RewriteStrategy.MULTI_QUERY,
num_queries=num_queries
)
def rewrite_contextual(
self,
query: str,
chat_history: List[Dict[str, str]]
) -> RewriteResult:
"""Convenience method for contextual rewriting."""
return self.rewrite(
query,
strategy=RewriteStrategy.CONTEXTUAL,
chat_history=chat_history
)
def add_abbreviation(self, abbrev: str, expansion: str) -> None:
"""Add a custom abbreviation expansion."""
self.abbreviations[abbrev.upper()] = expansion
def add_abbreviations(self, abbreviations: Dict[str, str]) -> None:
"""Add multiple custom abbreviation expansions."""
for abbrev, expansion in abbreviations.items():
self.abbreviations[abbrev.upper()] = expansion
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "praisonaiagents"
version = "0.0.165"
version = "0.0.168"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
requires-python = ">=3.10"
authors = [
@@ -0,0 +1,141 @@
"""
Query Rewriter Agent Example
Demonstrates various query rewriting strategies for improving RAG retrieval.
"""
from dotenv import load_dotenv
load_dotenv()
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
def main():
# Initialize the agent
agent = QueryRewriterAgent(
model="gpt-4o-mini",
verbose=True
)
print("=" * 60)
print("Query Rewriter Agent Examples")
print("=" * 60)
# 1. Basic Rewriting - Short keyword queries
print("\n1. BASIC REWRITING (Short Queries)")
print("-" * 40)
short_queries = [
"AI trends",
"python performance",
"RAG best practices"
]
for query in short_queries:
result = agent.rewrite(query, strategy=RewriteStrategy.BASIC)
print(f"Original: {query}")
print(f"Rewritten: {result.primary_query}")
print()
# 2. HyDE - Hypothetical Document Embeddings
print("\n2. HYDE (Hypothetical Document)")
print("-" * 40)
query = "What is quantum computing?"
result = agent.rewrite(query, strategy=RewriteStrategy.HYDE)
print(f"Original: {query}")
print("Hypothetical Document (truncated):")
print(f"{result.hypothetical_document[:300]}...")
print()
# 3. Step-Back - Higher-level concept questions
print("\n3. STEP-BACK (Broader Context)")
print("-" * 40)
query = "What is the difference between GPT-4 and Claude 3?"
result = agent.rewrite(query, strategy=RewriteStrategy.STEP_BACK)
print(f"Original: {query}")
print(f"Rewritten: {result.primary_query}")
print(f"Step-back Question: {result.step_back_question}")
print()
# 4. Sub-Queries - Decompose complex questions
print("\n4. SUB-QUERIES (Decomposition)")
print("-" * 40)
query = "How do I set up a RAG pipeline with vector search and what are the best embedding models to use?"
result = agent.rewrite(query, strategy=RewriteStrategy.SUB_QUERIES)
print(f"Original: {query}")
print("Sub-queries:")
for i, sq in enumerate(result.sub_queries, 1):
print(f" {i}. {sq}")
print()
# 5. Multi-Query - Multiple paraphrased versions
print("\n5. MULTI-QUERY (Ensemble Retrieval)")
print("-" * 40)
query = "How to improve LLM response quality?"
result = agent.rewrite(query, strategy=RewriteStrategy.MULTI_QUERY, num_queries=3)
print(f"Original: {query}")
print("Alternative queries:")
for i, q in enumerate(result.rewritten_queries, 1):
print(f" {i}. {q}")
print()
# 6. Contextual - Using conversation history
print("\n6. CONTEXTUAL (With Chat History)")
print("-" * 40)
chat_history = [
{"role": "user", "content": "Tell me about Python programming"},
{"role": "assistant", "content": "Python is a high-level programming language known for its simplicity and readability..."},
{"role": "user", "content": "What frameworks are popular?"},
{"role": "assistant", "content": "Popular Python frameworks include Django for web development, FastAPI for APIs, and PyTorch for ML..."}
]
query = "What about its performance compared to others?"
result = agent.rewrite(query, strategy=RewriteStrategy.CONTEXTUAL, chat_history=chat_history)
print(f"Original: {query}")
print(f"Standalone: {result.primary_query}")
print()
# 7. AUTO - Automatic strategy detection
print("\n7. AUTO (Automatic Strategy Detection)")
print("-" * 40)
test_queries = [
("ML", None), # Short query -> BASIC
("What about the cost?", chat_history), # Follow-up -> CONTEXTUAL
("Compare transformers vs RNNs and explain their use cases", None), # Complex -> SUB_QUERIES
]
for query, history in test_queries:
result = agent.rewrite(query, strategy=RewriteStrategy.AUTO, chat_history=history)
print(f"Query: {query}")
print(f"Detected Strategy: {result.strategy_used.value}")
print(f"Result: {result.primary_query}")
print()
# 8. Custom Abbreviations
print("\n8. CUSTOM ABBREVIATIONS")
print("-" * 40)
agent.add_abbreviations({
"K8s": "Kubernetes",
"TF": "TensorFlow",
"PT": "PyTorch"
})
query = "K8s deployment for TF models"
result = agent.rewrite(query, strategy=RewriteStrategy.BASIC)
print(f"Original: {query}")
print(f"Rewritten: {result.primary_query}")
print()
print("=" * 60)
print("Examples Complete!")
print("=" * 60)
if __name__ == "__main__":
main()
@@ -0,0 +1,403 @@
"""
Tests for QueryRewriterAgent
Run with: python -m pytest tests/query_rewriter_test.py -v
"""
import pytest
from unittest.mock import patch
class TestQueryRewriterImports:
"""Test that all imports work correctly."""
def test_import_from_package(self):
"""Test importing from main package."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy, RewriteResult
assert QueryRewriterAgent is not None
assert RewriteStrategy is not None
assert RewriteResult is not None
def test_import_from_agent_module(self):
"""Test importing from agent module."""
from praisonaiagents.agent import QueryRewriterAgent, RewriteStrategy, RewriteResult
assert QueryRewriterAgent is not None
assert RewriteStrategy is not None
assert RewriteResult is not None
class TestRewriteStrategy:
"""Test RewriteStrategy enum."""
def test_strategy_values(self):
"""Test all strategy values exist."""
from praisonaiagents import RewriteStrategy
assert RewriteStrategy.BASIC.value == "basic"
assert RewriteStrategy.HYDE.value == "hyde"
assert RewriteStrategy.STEP_BACK.value == "step_back"
assert RewriteStrategy.SUB_QUERIES.value == "sub_queries"
assert RewriteStrategy.MULTI_QUERY.value == "multi_query"
assert RewriteStrategy.CONTEXTUAL.value == "contextual"
assert RewriteStrategy.AUTO.value == "auto"
class TestRewriteResult:
"""Test RewriteResult dataclass."""
def test_result_creation(self):
"""Test creating a RewriteResult."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="test query",
rewritten_queries=["rewritten test query"],
strategy_used=RewriteStrategy.BASIC
)
assert result.original_query == "test query"
assert result.rewritten_queries == ["rewritten test query"]
assert result.strategy_used == RewriteStrategy.BASIC
def test_primary_query_property(self):
"""Test primary_query property."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="test",
rewritten_queries=["first", "second"],
strategy_used=RewriteStrategy.MULTI_QUERY
)
assert result.primary_query == "first"
def test_primary_query_fallback(self):
"""Test primary_query falls back to original if empty."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="test",
rewritten_queries=[],
strategy_used=RewriteStrategy.BASIC
)
assert result.primary_query == "test"
def test_all_queries_property(self):
"""Test all_queries property combines all queries."""
from praisonaiagents import RewriteResult, RewriteStrategy
result = RewriteResult(
original_query="original",
rewritten_queries=["rewritten"],
strategy_used=RewriteStrategy.STEP_BACK,
step_back_question="step back",
sub_queries=["sub1", "sub2"]
)
all_q = result.all_queries
assert "original" in all_q
assert "rewritten" in all_q
assert "step back" in all_q
assert "sub1" in all_q
assert "sub2" in all_q
class TestQueryRewriterAgentInit:
"""Test QueryRewriterAgent initialization."""
def test_default_initialization(self):
"""Test agent initializes with defaults."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent.name == "QueryRewriterAgent"
assert agent.model == "gpt-4o-mini"
assert not agent.verbose
assert agent.max_queries == 5
assert agent.temperature == 0.3
assert agent.max_tokens == 500
def test_custom_initialization(self):
"""Test agent initializes with custom values."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent(
name="CustomRewriter",
model="gpt-4o",
verbose=True,
max_queries=10,
temperature=0.5,
max_tokens=1000
)
assert agent.name == "CustomRewriter"
assert agent.model == "gpt-4o"
assert agent.verbose
assert agent.max_queries == 10
assert agent.temperature == 0.5
assert agent.max_tokens == 1000
def test_default_abbreviations(self):
"""Test default abbreviations are set."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert "AI" in agent.abbreviations
assert "ML" in agent.abbreviations
assert "RAG" in agent.abbreviations
assert agent.abbreviations["AI"] == "Artificial Intelligence"
class TestAbbreviationExpansion:
"""Test abbreviation expansion functionality."""
def test_expand_abbreviations(self):
"""Test _expand_abbreviations method."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._expand_abbreviations("AI trends")
assert "Artificial Intelligence" in result
def test_add_abbreviation(self):
"""Test adding single abbreviation."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
agent.add_abbreviation("K8s", "Kubernetes")
assert "K8S" in agent.abbreviations
assert agent.abbreviations["K8S"] == "Kubernetes"
def test_add_abbreviations(self):
"""Test adding multiple abbreviations."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
agent.add_abbreviations({
"TF": "TensorFlow",
"PT": "PyTorch"
})
assert "TF" in agent.abbreviations
assert "PT" in agent.abbreviations
class TestQueryDetection:
"""Test query analysis methods."""
def test_is_short_query(self):
"""Test short query detection."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent._is_short_query("AI")
assert agent._is_short_query("AI trends")
assert not agent._is_short_query("What are the latest AI trends in 2025?")
def test_is_follow_up(self):
"""Test follow-up query detection."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent._is_follow_up("What about it?")
assert agent._is_follow_up("And this one?")
assert not agent._is_follow_up("What is AI?")
def test_is_complex_query(self):
"""Test complex query detection."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert agent._is_complex_query("What is A and what is B?")
assert agent._is_complex_query("Question one? Question two?")
assert not agent._is_complex_query("What is AI?")
class TestStrategyDetection:
"""Test automatic strategy detection."""
def test_detect_basic_for_short(self):
"""Test BASIC detected for short queries."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent()
strategy = agent._detect_strategy("AI")
assert strategy == RewriteStrategy.BASIC
def test_detect_contextual_for_followup(self):
"""Test CONTEXTUAL detected for follow-up with history."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent()
chat_history = [{"role": "user", "content": "Tell me about Python"}]
strategy = agent._detect_strategy("What about it?", chat_history)
assert strategy == RewriteStrategy.CONTEXTUAL
def test_detect_sub_queries_for_complex(self):
"""Test SUB_QUERIES detected for complex queries."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
agent = QueryRewriterAgent()
strategy = agent._detect_strategy("What is A and what is B?")
assert strategy == RewriteStrategy.SUB_QUERIES
class TestJSONParsing:
"""Test JSON array parsing from LLM responses."""
def test_parse_valid_json(self):
"""Test parsing valid JSON array."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._parse_json_array('["query1", "query2"]')
assert result == ["query1", "query2"]
def test_parse_json_in_text(self):
"""Test extracting JSON from surrounding text."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._parse_json_array('Here are the queries: ["q1", "q2"]')
assert result == ["q1", "q2"]
def test_parse_fallback_newlines(self):
"""Test fallback to newline splitting."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
result = agent._parse_json_array("1. Query one\n2. Query two\n3. Query three")
assert len(result) >= 2
class TestRewriteWithMockedLLM:
"""Test rewrite methods with mocked LLM."""
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_basic(self, mock_llm):
"""Test basic rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = "What are the current trends in AI?"
agent = QueryRewriterAgent()
result = agent.rewrite("AI trends", strategy=RewriteStrategy.BASIC)
assert result.strategy_used == RewriteStrategy.BASIC
assert len(result.rewritten_queries) > 0
mock_llm.assert_called_once()
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_hyde(self, mock_llm):
"""Test HyDE rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = "Quantum computing is a type of computation..."
agent = QueryRewriterAgent()
result = agent.rewrite("What is quantum computing?", strategy=RewriteStrategy.HYDE)
assert result.strategy_used == RewriteStrategy.HYDE
assert result.hypothetical_document is not None
mock_llm.assert_called_once()
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_step_back(self, mock_llm):
"""Test step-back rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
# First call for step-back, second for basic
mock_llm.side_effect = [
"What are the fundamentals of language models?",
"What are the differences between GPT-4 and Claude 3?"
]
agent = QueryRewriterAgent()
result = agent.rewrite("GPT-4 vs Claude 3?", strategy=RewriteStrategy.STEP_BACK)
assert result.strategy_used == RewriteStrategy.STEP_BACK
assert result.step_back_question is not None
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_sub_queries(self, mock_llm):
"""Test sub-query decomposition with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = '["What is RAG?", "What are the best embedding models?"]'
agent = QueryRewriterAgent()
result = agent.rewrite("RAG setup and embedding models?", strategy=RewriteStrategy.SUB_QUERIES)
assert result.strategy_used == RewriteStrategy.SUB_QUERIES
assert result.sub_queries is not None
assert len(result.sub_queries) >= 1
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_multi_query(self, mock_llm):
"""Test multi-query generation with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = '["Query 1", "Query 2", "Query 3"]'
agent = QueryRewriterAgent()
result = agent.rewrite("LLM quality?", strategy=RewriteStrategy.MULTI_QUERY, num_queries=3)
assert result.strategy_used == RewriteStrategy.MULTI_QUERY
assert len(result.rewritten_queries) >= 1
@patch.object(__import__('praisonaiagents.agent.query_rewriter_agent', fromlist=['QueryRewriterAgent']).QueryRewriterAgent, '_call_agent')
def test_rewrite_contextual(self, mock_llm):
"""Test contextual rewriting with mocked LLM."""
from praisonaiagents import QueryRewriterAgent, RewriteStrategy
mock_llm.return_value = "How does Python's performance compare to other languages?"
agent = QueryRewriterAgent()
chat_history = [
{"role": "user", "content": "Tell me about Python"},
{"role": "assistant", "content": "Python is a programming language..."}
]
result = agent.rewrite(
"What about its performance?",
strategy=RewriteStrategy.CONTEXTUAL,
chat_history=chat_history
)
assert result.strategy_used == RewriteStrategy.CONTEXTUAL
assert "Python" in result.primary_query or len(result.rewritten_queries) > 0
class TestConvenienceMethods:
"""Test convenience methods."""
def test_convenience_methods_exist(self):
"""Test all convenience methods exist."""
from praisonaiagents import QueryRewriterAgent
agent = QueryRewriterAgent()
assert hasattr(agent, 'rewrite_basic')
assert hasattr(agent, 'rewrite_hyde')
assert hasattr(agent, 'rewrite_step_back')
assert hasattr(agent, 'rewrite_sub_queries')
assert hasattr(agent, 'rewrite_multi_query')
assert hasattr(agent, 'rewrite_contextual')
if __name__ == "__main__":
pytest.main([__file__, "-v"])
+1 -1
View File
@@ -2422,7 +2422,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.0.165"
version = "0.0.168"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },
+29 -4
View File
@@ -1,9 +1,34 @@
# PraisonAI Package
This is the PraisonAI package, which serves as a wrapper for PraisonAIAgents.
## Development
It provides a simple and intuitive interface for working with AI agents and their capabilities.
```bash
# Install
uv sync
## Directory Structure
# Build
uv build
The main package code is located in the `praisonai` subdirectory.
# Publish
uv publish
```
## Version Bump
```bash
# From project root
python src/praisonai/scripts/bump_version.py 2.2.96
# With praisonaiagents dependency
python src/praisonai/scripts/bump_version.py 2.2.96 --agents 0.0.167
```
## Release
```bash
# From project root (runs uv lock, build, git tag, gh release)
python src/praisonai/scripts/release.py
# Then publish
cd src/praisonai && uv publish
```
+2 -2
View File
@@ -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.93.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.2.93.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]"
+191 -9
View File
@@ -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,6 +577,9 @@ 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 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/)")
parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose output for research")
@@ -721,7 +723,10 @@ class PraisonAI:
research_model = getattr(args, 'model', None)
verbose = getattr(args, 'verbose', False)
save = getattr(args, 'save', False)
self.handle_research_command(research_query, research_model, verbose, save)
query_rewrite = getattr(args, 'query_rewrite', False)
rewrite_tools = getattr(args, 'rewrite_tools', None)
tools_path = getattr(args, 'tools', None)
self.handle_research_command(research_query, research_model, verbose, save, query_rewrite, tools_path, rewrite_tools)
sys.exit(0)
# Only check framework availability for agent-related operations
@@ -742,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",
@@ -961,7 +1061,7 @@ class PraisonAI:
print(f"[red]ERROR: Context engineering failed: {e}[/red]")
sys.exit(1)
def handle_research_command(self, query: str, model: str = None, verbose: bool = False, save: bool = False) -> str:
def handle_research_command(self, query: str, model: str = None, verbose: bool = False, save: bool = False, query_rewrite: bool = False, tools_path: str = None, rewrite_tools: str = None) -> str:
"""
Handle the research command by creating a DeepResearchAgent and running it.
@@ -970,6 +1070,9 @@ class PraisonAI:
model: Model for deep research (optional, defaults to o4-mini-deep-research)
verbose: Enable verbose output (default: False)
save: Save output to file (default: False)
query_rewrite: Rewrite query before research (default: False)
tools_path: Path to tools.py file with custom tools (default: None)
rewrite_tools: Tools for query rewriter (tool names or file path)
Returns:
str: Research report
@@ -984,6 +1087,10 @@ class PraisonAI:
logging.getLogger('httpx').setLevel(logging.WARNING)
logging.getLogger('httpcore').setLevel(logging.WARNING)
# Rewrite query if requested
if query_rewrite:
query = self._rewrite_query(query, rewrite_tools, verbose)
print("[bold green]Starting Deep Research...[/bold green]")
print(f"Query: {query}")
@@ -993,11 +1100,86 @@ class PraisonAI:
print(f"Model: {model}")
# Create DeepResearchAgent (verbose=True is default for streaming output)
agent = DeepResearchAgent(model=model)
# Load tools if specified
tools_list = []
if tools_path:
# Check if it's a file path or comma-separated tool names
if os.path.isfile(tools_path):
# Load from file
try:
import inspect
spec = importlib.util.spec_from_file_location("tools_module", tools_path)
if spec and spec.loader:
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
# Get all callable functions from the module
for name, obj in inspect.getmembers(module):
if inspect.isfunction(obj) and not name.startswith('_'):
tools_list.append(obj)
if tools_list:
print(f"[cyan]Loaded {len(tools_list)} tools from {tools_path}[/cyan]")
except Exception as e:
print(f"[yellow]Warning: Failed to load tools from {tools_path}: {e}[/yellow]")
else:
# Treat as comma-separated tool names (e.g., "internet_search,wiki_search")
try:
from praisonaiagents.tools import TOOL_MAPPINGS
import praisonaiagents.tools as tools_module
tool_names = [t.strip() for t in tools_path.split(',')]
for tool_name in tool_names:
if tool_name in TOOL_MAPPINGS:
try:
tool = getattr(tools_module, tool_name)
tools_list.append(tool)
except Exception as e:
print(f"[yellow]Warning: Failed to load tool '{tool_name}': {e}[/yellow]")
else:
print(f"[yellow]Warning: Unknown tool '{tool_name}'[/yellow]")
if tools_list:
print(f"[cyan]Loaded {len(tools_list)} built-in tools: {', '.join(tool_names)}[/cyan]")
except ImportError:
print("[yellow]Warning: Could not import tools module[/yellow]")
# Execute the research
result = agent.research(query)
# If tools are provided, use Agent with tools first, then DeepResearchAgent
if tools_list:
from praisonaiagents import Agent, Task, PraisonAIAgents
# Create a research assistant agent with tools
research_assistant = Agent(
name="Research Assistant",
role="Information Gatherer",
goal="Gather relevant information using available tools",
backstory="You are an expert at using tools to gather information for research.",
tools=tools_list,
llm="gpt-4o-mini",
verbose=False
)
# Create task to gather initial information
gather_task = Task(
description=f"Use your tools to gather relevant information about: {query}",
expected_output="A summary of information gathered from the tools",
agent=research_assistant
)
print("[cyan]Gathering information with tools...[/cyan]")
agents = PraisonAIAgents(agents=[research_assistant], tasks=[gather_task], verbose=0)
tool_results = agents.start()
# Enhance query with tool results
enhanced_query = f"{query}\n\nAdditional context from tools:\n{tool_results}"
print("[cyan]Tools context gathered, starting deep research...[/cyan]")
# Create DeepResearchAgent
agent = DeepResearchAgent(model=model)
result = agent.research(enhanced_query)
else:
# Create DeepResearchAgent (verbose=True is default for streaming output)
agent = DeepResearchAgent(model=model)
# Execute the research
result = agent.research(query)
print("\n[bold green]Research Complete![/bold green]")
print("\n" + "="*60)
+1 -1
View File
@@ -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.93 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
View File
@@ -1 +1 @@
__version__ = "0.0.1"
__version__ = "2.2.97"
+11 -194
View File
@@ -1,9 +1,9 @@
[project]
name = "PraisonAI"
version = "2.2.93"
dynamic = ["version"]
description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human-agent collaboration."
readme = "README.md"
license = ""
license = {text = "MIT"}
requires-python = ">=3.10"
authors = [
{ name = "Mervin Praison" }
@@ -12,7 +12,7 @@ dependencies = [
"rich>=13.7",
"markdown>=3.5",
"pyparsing>=3.0.0",
"praisonaiagents>=0.0.165",
"praisonaiagents>=0.0.168",
"python-dotenv>=0.19.0",
"instructor>=1.3.3",
"PyYAML>=6.0",
@@ -100,200 +100,17 @@ autogen-v4 = [
"crewai"
]
[tool.poetry]
name = "PraisonAI"
version = "2.2.93"
description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human-agent collaboration."
authors = ["Mervin Praison"]
license = ""
readme = "README.md"
[tool.poetry.urls]
[project.urls]
Homepage = "https://docs.praison.ai"
Repository = "https://github.com/mervinpraison/PraisonAI"
[tool.setuptools]
packages = ["praisonai"]
[tool.poetry.dependencies]
python = ">=3.10,<3.13"
rich = ">=13.7"
markdown = ">=3.5"
pyparsing = ">=3.0.0"
praisonaiagents = ">=0.0.165"
python-dotenv = ">=0.19.0"
instructor = ">=1.3.3"
PyYAML = ">=6.0"
mcp = ">=1.6.0"
pyautogen = {version = "==0.2.29", optional = true}
autogen-agentchat = {version = ">=0.4.0", optional = true}
autogen-ext = {version = ">=0.4.0", optional = true}
autogen-core = {version = ">=0.4.0", optional = true}
crewai = {version = ">=0.157.0", optional = true}
praisonai-tools = {version = ">=0.0.22", optional = true}
chainlit = {version = "==2.5.5", optional = true}
gradio = {version = ">=4.26.0", optional = true}
flask = {version = ">=3.0.0", optional = true}
agentops = {version = ">=0.3.12", optional = true}
langchain-google-genai = {version = ">=2.1.0", optional = true}
langchain-anthropic = {version = ">=0.3.0", optional = true}
langchain-openai = {version = ">=0.2.1,<0.3.0", optional = true}
langchain-cohere = {version = ">=0.3.0,<0.4.0", optional = true}
litellm = {version = ">=1.72.6", optional = true}
aiosqlite= {version = ">=0.20.0", optional = true}
greenlet = {version = ">=3.0.3", optional = true}
tavily-python = {version = "==0.5.0", optional=true}
crawl4ai = {version = ">=0.7.0", optional = true}
websockets = {version = ">=12.0", optional = true}
plotly = {version = ">=5.24.0", optional = true}
yfinance = {version = ">=0.2.44", optional = true}
duckduckgo_search = {version = ">=6.3.0", optional = true}
twilio = {version = ">=7.0.0", optional = true}
fastapi = {version = ">=0.95.0", optional = true}
uvicorn = {version = ">=0.20.0", optional = true}
flaml = {version = ">=2.3.1", extras = ["automl"], optional = true}
pyngrok = {version = ">=1.4.0", optional = true}
sqlalchemy = {version = ">=2.0.36", optional = true}
playwright = {version = ">=1.47.0", optional = true}
openai = {version = ">=1.54.0", optional = true}
pydantic = {version = "<=2.10.1", optional = true}
# unsloth = {version = ">=2024.11.7", extras = ["colab-new"], optional = true}
[tool.poetry.group.docs.dependencies]
mkdocs = "*"
markdown-include = "*"
mkdocs-material = "*"
mkdocs-jupyter = "*"
mkdocs-glightbox = "*"
pymdown-extensions = "*"
Pygments = "*"
mkdocstrings = "*"
mkdocs-apidoc = "*"
mkdocstrings-python = "*"
pdoc3 = "*"
[tool.poetry.group.test.dependencies]
pytest = "8.2.2"
pytest-asyncio = ">=0.26.0"
pre-commit = "3.7.1"
unittest-xml-reporting = "3.2.0"
xmlrunner = "*"
unittest2 = "*"
[tool.poetry.group.dev.dependencies]
pytest = "8.2.2"
pytest-asyncio = ">=0.26.0"
pre-commit = "3.7.1"
unittest-xml-reporting = "3.2.0"
mkdocs = "*"
markdown-include = "*"
mkdocs-material = "*"
mkdocs-jupyter = "*"
mkdocs-glightbox = "*"
pymdown-extensions = "*"
Pygments = "*"
mkdocstrings = "*"
mkdocs-apidoc = "*"
mkdocstrings-python = "*"
pdoc3 = "*"
xmlrunner = "*"
unittest2 = "*"
chainlit = "==2.5.5"
gradio = ">=4.26.0"
flask = ">=3.0.0"
agentops = ">=0.3.12"
langchain-google-genai = ">=2.1.0"
langchain-anthropic = ">=0.3.0"
langchain-openai = ">=0.2.1,<0.3.0"
langchain-cohere = ">=0.3.0,<0.4.0"
litellm = ">=1.72.6"
aiosqlite = ">=0.20.0"
greenlet = ">=3.0.3"
[build-system]
requires = [
"poetry-core",
"setuptools>=42",
]
build-backend = "poetry.core.masonry.api"
requires = ["setuptools>=64"]
build-backend = "setuptools.build_meta"
[tool.setuptools.dynamic]
version = {attr = "praisonai.version.__version__"}
[tool.poetry.scripts]
praisonai = "praisonai.__main__:main"
setup-conda-env = "setup.setup_conda_env:main"
praisonai-call = "praisonai.api.call:main"
[tool.poetry.extras]
ui = [
"chainlit",
"aiosqlite",
"greenlet",
"sqlalchemy",
"pydantic"
]
gradio = ["gradio"]
api = ["flask"]
agentops = ["agentops"]
google = ["langchain-google-genai"]
openai = ["langchain-openai"]
anthropic = ["langchain-anthropic"]
cohere = ["langchain-cohere"]
chat = [
"chainlit",
"litellm",
"aiosqlite",
"greenlet",
"tavily-python",
"crawl4ai",
"sqlalchemy",
"playwright",
"rich",
"pydantic"
]
code = [
"chainlit",
"litellm",
"aiosqlite",
"greenlet",
"tavily-python",
"crawl4ai",
"sqlalchemy",
"playwright",
"pydantic"
]
train = []
realtime = [
"chainlit",
"litellm",
"aiosqlite",
"greenlet",
"tavily-python",
"crawl4ai",
"websockets",
"plotly",
"yfinance",
"duckduckgo_search",
"sqlalchemy"
]
call = [
"twilio",
"fastapi",
"uvicorn",
"websockets",
"flaml",
"pyngrok",
"rich",
"openai",
]
crewai = ["crewai", "praisonai-tools"]
autogen = ["pyautogen", "praisonai-tools", "crewai"]
autogen-v4 = ["autogen-agentchat", "autogen-ext", "autogen-core", "praisonai-tools", "crewai"]
[tool.poetry-dynamic-versioning]
enable = true
vcs = "git"
style = "semver"
[tool.poetry.build]
generate-setup-file = false
script = "praisonai/setup/post_install.py"
[tool.setuptools.packages.find]
where = ["."]
include = ["praisonai*"]
+142
View File
@@ -0,0 +1,142 @@
#!/usr/bin/env python3
"""
Version bump script for PraisonAI package.
This script updates the version number in all required locations:
- praisonai/version.py (single source of truth for Python package)
- praisonai/deploy.py (Dockerfile template)
- ../../docker/Dockerfile, Dockerfile.chat, Dockerfile.dev, Dockerfile.ui
- praisonai.rb (Homebrew formula)
Usage:
python src/praisonai/scripts/bump_version.py 2.2.96
Or with optional praisonaiagents version:
python src/praisonai/scripts/bump_version.py 2.2.96 --agents 0.0.167
"""
import re
import sys
import argparse
from pathlib import Path
def get_project_root() -> Path:
"""Get the project root directory (praisonai-package)."""
# scripts/ -> src/praisonai/ -> src/ -> praisonai-package/
return Path(__file__).parent.parent.parent.parent
def update_file(filepath: Path, patterns: list[tuple[str, str]]) -> bool:
"""Update a file with the given regex patterns and replacements."""
if not filepath.exists():
print(f" ⚠️ File not found: {filepath}")
return False
content = filepath.read_text()
original = content
for pattern, replacement in patterns:
content = re.sub(pattern, replacement, content)
if content != original:
filepath.write_text(content)
print(f" ✅ Updated: {filepath.relative_to(get_project_root())}")
return True
else:
print(f" ⏭️ No changes: {filepath.relative_to(get_project_root())}")
return False
def bump_version(new_version: str, agents_version: str | None = None):
"""Bump version in all required files."""
root = get_project_root()
print(f"\n🚀 Bumping PraisonAI version to {new_version}\n")
praisonai_dir = root / "src/praisonai"
# 1. Update version.py (single source of truth)
print("📦 Python Package:")
update_file(
praisonai_dir / "praisonai/version.py",
[(r'__version__ = "[^"]+"', f'__version__ = "{new_version}"')]
)
# 2. Update deploy.py (Dockerfile template)
print("\n🐳 Deploy Script:")
update_file(
praisonai_dir / "praisonai/deploy.py",
[(r'praisonai==[0-9.]+', f'praisonai=={new_version}')]
)
# 3. Update Dockerfiles
print("\n🐳 Dockerfiles:")
dockerfiles = [
"docker/Dockerfile",
"docker/Dockerfile.chat",
"docker/Dockerfile.dev",
"docker/Dockerfile.ui",
]
for dockerfile in dockerfiles:
update_file(
root / dockerfile,
[(r'"praisonai>=[0-9.]+"', f'"praisonai>={new_version}"')]
)
# 4. Update Homebrew formula
print("\n🍺 Homebrew Formula:")
update_file(
praisonai_dir / "praisonai.rb",
[(r'v[0-9]+\.[0-9]+\.[0-9]+', f'v{new_version}')]
)
# 5. Update praisonaiagents dependency if specified
if agents_version:
print(f"\n📦 Updating praisonaiagents dependency to {agents_version}:")
update_file(
praisonai_dir / "pyproject.toml",
[(r'praisonaiagents>=[0-9.]+', f'praisonaiagents>={agents_version}')]
)
print("\n✨ Version bump complete!")
print("\nNext steps:")
print(" 1. Run: cd src/praisonai && uv lock")
print(" 2. Run: cd src/praisonai && uv build")
print(f" 3. Commit changes: git add -A && git commit -m 'Bump version to {new_version}'")
print(f" 4. Tag release: git tag v{new_version}")
print(" 5. Publish: cd src/praisonai && uv publish")
def main():
parser = argparse.ArgumentParser(
description="Bump PraisonAI version in all required files"
)
parser.add_argument(
"version",
help="New version number (e.g., 2.2.96)"
)
parser.add_argument(
"--agents", "-a",
help="Optional: Update praisonaiagents dependency version (e.g., 0.0.167)",
default=None
)
args = parser.parse_args()
# Validate version format
if not re.match(r'^\d+\.\d+\.\d+$', args.version):
print(f"❌ Invalid version format: {args.version}")
print(" Expected format: X.Y.Z (e.g., 2.2.96)")
sys.exit(1)
if args.agents and not re.match(r'^\d+\.\d+\.\d+$', args.agents):
print(f"❌ Invalid agents version format: {args.agents}")
print(" Expected format: X.Y.Z (e.g., 0.0.167)")
sys.exit(1)
bump_version(args.version, args.agents)
if __name__ == "__main__":
main()
+102
View File
@@ -0,0 +1,102 @@
#!/usr/bin/env python3
"""
Release script for PraisonAI package.
Usage:
python src/praisonai/scripts/release.py
This script:
1. Reads version from praisonai/version.py
2. Runs uv lock
3. Runs uv build
4. Commits changes
5. Creates git tag
6. Pushes to GitHub
7. Creates GitHub release (latest)
"""
import subprocess
import sys
from pathlib import Path
def get_project_root() -> Path:
"""Get the project root directory (praisonai-package)."""
return Path(__file__).parent.parent.parent.parent
def get_praisonai_dir() -> Path:
"""Get the praisonai package directory."""
return get_project_root() / "src/praisonai"
def get_version() -> str:
"""Read version from version.py."""
version_file = get_praisonai_dir() / "praisonai/version.py"
content = version_file.read_text()
for line in content.splitlines():
if line.startswith("__version__"):
return line.split('"')[1]
raise ValueError("Could not find version in version.py")
def run(cmd: list[str], cwd: Path | None = None, check: bool = True) -> subprocess.CompletedProcess:
"""Run a command and print it."""
print(f"$ {' '.join(cmd)}")
result = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
if result.stdout:
print(result.stdout)
if result.stderr:
print(result.stderr, file=sys.stderr)
if check and result.returncode != 0:
print(f"❌ Command failed with exit code {result.returncode}")
sys.exit(1)
return result
def main():
root = get_project_root()
praisonai_dir = get_praisonai_dir()
version = get_version()
tag = f"v{version}"
print(f"\n🚀 Releasing PraisonAI {tag}\n")
# 1. uv lock
print("📦 Running uv lock...")
run(["uv", "lock"], cwd=praisonai_dir)
# 2. uv build
print("\n🔨 Running uv build...")
run(["uv", "build"], cwd=praisonai_dir)
# 3. Git add and commit
print("\n📝 Committing changes...")
run(["git", "add", "-A"], cwd=root)
run(["git", "commit", "-m", f"Release {tag}"], cwd=root, check=False)
# 4. Create git tag
print(f"\n🏷️ Creating tag {tag}...")
run(["git", "tag", "-f", tag], cwd=root)
# 5. Push to GitHub
print("\n⬆️ Pushing to GitHub...")
run(["git", "push"], cwd=root)
run(["git", "push", "--tags", "-f"], cwd=root)
# 6. Create GitHub release
print(f"\n🎉 Creating GitHub release {tag}...")
run([
"gh", "release", "create", tag,
"--title", f"PraisonAI {tag}",
"--notes", f"Release {tag}",
"--latest"
], cwd=root)
print(f"\n✅ Released PraisonAI {tag}")
print("\nNext step:")
print(" cd src/praisonai && uv publish")
if __name__ == "__main__":
main()
@@ -75,7 +75,8 @@ roles:
def test_autogen_environment_check(self):
"""Verify AutoGen environment is properly configured"""
# Check API key is available
assert os.getenv("OPENAI_API_KEY"), "OPENAI_API_KEY required for real tests"
if not os.getenv("OPENAI_API_KEY"):
pytest.skip("OPENAI_API_KEY required for real tests")
# Check AutoGen can be imported
try:
@@ -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"])
+4 -5
View File
@@ -4044,7 +4044,6 @@ wheels = [
[[package]]
name = "praisonai"
version = "2.2.93"
source = { editable = "." }
dependencies = [
{ name = "instructor" },
@@ -4198,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.165" },
{ 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" },
@@ -4253,7 +4252,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.0.165"
version = "0.0.168"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiohttp" },
@@ -4262,9 +4261,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "rich" },
]
sdist = { url = "https://files.pythonhosted.org/packages/c3/49/37f8518b2205760cec66ac44482c1cb9e33c0cb568c4c6cb07f9908dceb3/praisonaiagents-0.0.165.tar.gz", hash = "sha256:b2a43a44d51454f9917a2345407bf86fdc46046bafac96601c98d1d21fb2601a", size = 298269, upload-time = "2025-12-13T02:18:34.883Z" }
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/26/e7/c3b21f4c68629612ab4e4d128fca5fdabe141df16b0861679e2736adebe7/praisonaiagents-0.0.165-py3-none-any.whl", hash = "sha256:65b880f6d4c5d05fadf86153338306e12004316da9dc9f83942b9e16f2133851", size = 315370, upload-time = "2025-12-13T02:18:33.197Z" },
{ 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]]