mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-08 01:08:39 -10:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e29570969d |
No files matched your search
@@ -204,6 +204,12 @@ npm install praisonai
|
||||
| ↳ Add Custom Knowledge | [Example](examples/python/concepts/knowledge-agents.py) | [📖](https://docs.praison.ai/features/knowledge) |
|
||||
| ↳ RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [📖](https://docs.praison.ai/features/rag) |
|
||||
| ↳ Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [📖](https://docs.praison.ai/features/chat-with-pdf) |
|
||||
| ↳ Data Readers (PDF, DOCX, etc.) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-readers-api) |
|
||||
| ↳ Vector Store Selection | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-vector-store-api) |
|
||||
| ↳ Retrieval Strategies | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-retrieval-api) |
|
||||
| ↳ Rerankers | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-reranker-api) |
|
||||
| ↳ Index Types (Vector/Keyword/Hybrid) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-index-api) |
|
||||
| ↳ Query Engines (Sub-Question, etc.) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-query-engine-api) |
|
||||
| **🔬 Research & Intelligence** | | |
|
||||
| ↳ Deep Research Agents | [Example](examples/python/agents/research-agent.py) | [📖](https://docs.praison.ai/agents/deep-research) |
|
||||
| ↳ Query Rewriter Agent | [Example](#5-query-rewriter-agent) | [📖](https://docs.praison.ai/agents/query-rewriter) |
|
||||
@@ -2641,6 +2647,84 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
|
||||
---
|
||||
|
||||
## 📚 Knowledge & Retrieval (RAG)
|
||||
|
||||
PraisonAI provides a complete knowledge stack for building RAG applications with multiple vector stores, retrieval strategies, rerankers, and query modes.
|
||||
|
||||
### Knowledge CLI Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai knowledge add <file\|dir\|url>` | Add documents to knowledge base |
|
||||
| `praisonai knowledge query <question>` | Query knowledge base with RAG |
|
||||
| `praisonai knowledge list` | List indexed documents |
|
||||
| `praisonai knowledge clear` | Clear knowledge base |
|
||||
| `praisonai knowledge stats` | Show knowledge base statistics |
|
||||
|
||||
### Knowledge CLI Options
|
||||
|
||||
| Option | Values | Description |
|
||||
|--------|--------|-------------|
|
||||
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | Vector store backend |
|
||||
| `--retrieval` | `basic`, `fusion`, `recursive`, `auto_merge` | Retrieval strategy |
|
||||
| `--reranker` | `simple`, `llm`, `cross_encoder`, `cohere` | Reranking method |
|
||||
| `--index-type` | `vector`, `keyword`, `hybrid` | Index type |
|
||||
| `--query-mode` | `default`, `sub_question`, `summarize` | Query mode |
|
||||
|
||||
### Knowledge CLI Examples
|
||||
|
||||
```bash
|
||||
# Add documents
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge add https://example.com/page.html
|
||||
praisonai knowledge add "*.pdf"
|
||||
|
||||
# Query with advanced options
|
||||
praisonai knowledge query "How to authenticate?" --retrieval fusion --reranker llm
|
||||
|
||||
# Full advanced query
|
||||
praisonai knowledge query "authentication flow" \
|
||||
--vector-store chroma \
|
||||
--retrieval fusion \
|
||||
--reranker llm \
|
||||
--index-type hybrid \
|
||||
--query-mode sub_question
|
||||
```
|
||||
|
||||
### Knowledge SDK Usage
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, Knowledge
|
||||
|
||||
# Simple usage with Agent
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
knowledge=["docs/manual.pdf", "data/faq.txt"],
|
||||
knowledge_config={
|
||||
"vector_store": {"provider": "chroma"}
|
||||
}
|
||||
)
|
||||
response = agent.chat("How do I authenticate?")
|
||||
|
||||
# Direct Knowledge usage
|
||||
knowledge = Knowledge()
|
||||
knowledge.add("document.pdf")
|
||||
results = knowledge.search("authentication", limit=5)
|
||||
```
|
||||
|
||||
### Knowledge Stack Features Table
|
||||
|
||||
| Feature | Description | SDK Docs | CLI Docs |
|
||||
|---------|-------------|----------|----------|
|
||||
| **Data Readers** | Load PDF, Markdown, Text, HTML, URLs | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Vector Stores** | ChromaDB, Pinecone, Qdrant, Weaviate, In-Memory | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Retrieval Strategies** | Basic, Fusion (RRF), Recursive, Auto-Merge | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Rerankers** | Simple, LLM, Cross-Encoder, Cohere | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Index Types** | Vector, Keyword (BM25), Hybrid | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Query Engines** | Default, Sub-Question, Summarize | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
|
||||
---
|
||||
|
||||
## 🔬 Advanced Features
|
||||
|
||||
### Research & Intelligence
|
||||
@@ -2755,6 +2839,44 @@ agent.chat("Hello!") # Auto-persists messages, runs, traces
|
||||
| `praisonai persistence migrate` | Apply schema migrations |
|
||||
| `praisonai persistence status` | Show schema status |
|
||||
|
||||
### Knowledge CLI Commands {#knowledge-cli}
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai knowledge add <source>` | Add file, directory, URL, or glob pattern |
|
||||
| `praisonai knowledge query "<question>"` | Query knowledge base with RAG |
|
||||
| `praisonai knowledge list` | List indexed documents |
|
||||
| `praisonai knowledge clear` | Clear knowledge base |
|
||||
| `praisonai knowledge stats` | Show knowledge base statistics |
|
||||
|
||||
**Knowledge Query Flags:**
|
||||
|
||||
| Flag | Values | Default |
|
||||
|------|--------|---------|
|
||||
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | `chroma` |
|
||||
| `--retrieval-strategy` | `basic`, `fusion`, `recursive`, `auto_merge` | `basic` |
|
||||
| `--reranker` | `none`, `simple`, `llm`, `cross_encoder`, `cohere` | `none` |
|
||||
| `--index-type` | `vector`, `keyword`, `hybrid` | `vector` |
|
||||
| `--query-mode` | `default`, `sub_question`, `summarize` | `default` |
|
||||
| `--workspace` | Path to workspace directory | Current dir |
|
||||
| `--session` | Session ID for persistence | - |
|
||||
|
||||
**Examples:**
|
||||
|
||||
```bash
|
||||
# Add documents
|
||||
praisonai knowledge add document.pdf
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge add "*.md"
|
||||
|
||||
# Query with options
|
||||
praisonai knowledge query "How to authenticate?" \
|
||||
--vector-store chroma \
|
||||
--retrieval-strategy fusion \
|
||||
--reranker simple \
|
||||
--query-mode sub_question
|
||||
```
|
||||
|
||||
### Databases Table
|
||||
|
||||
| Database | Store Type | Install | Example | Docs |
|
||||
|
||||
@@ -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.5.0" \
|
||||
"praisonai>=2.5.1" \
|
||||
"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.5.0" \
|
||||
"praisonai>=2.5.1" \
|
||||
"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.5.0" \
|
||||
"praisonai>=2.5.1" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
Async Runs Examples
|
||||
===================
|
||||
|
||||
This directory contains examples for the PraisonAI Async Jobs API.
|
||||
|
||||
Examples:
|
||||
---------
|
||||
1. sdk_submit_poll_result.py - Python SDK example for submit/poll/result
|
||||
2. cli_submit_poll_result.sh - Bash/curl example for submit/poll/result
|
||||
3. idempotency.sh - Demonstrates idempotency key usage
|
||||
4. cancel.sh - Demonstrates job cancellation
|
||||
5. sse_stream.sh - Demonstrates SSE streaming
|
||||
|
||||
Prerequisites:
|
||||
--------------
|
||||
1. Set OPENAI_API_KEY environment variable
|
||||
2. Start the jobs server:
|
||||
|
||||
cd /path/to/praisonai-package/src/praisonai
|
||||
python -m uvicorn praisonai.jobs.server:create_app --port 8005 --factory
|
||||
|
||||
3. Run examples:
|
||||
|
||||
python sdk_submit_poll_result.py
|
||||
bash cli_submit_poll_result.sh
|
||||
bash idempotency.sh
|
||||
bash cancel.sh
|
||||
bash sse_stream.sh
|
||||
|
||||
Environment Variables:
|
||||
----------------------
|
||||
- OPENAI_API_KEY: Required for agent execution
|
||||
- PRAISONAI_BASE_URL: API server URL (default: http://127.0.0.1:8005)
|
||||
"""
|
||||
@@ -0,0 +1,54 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: Cancel a Running Job
|
||||
#
|
||||
# This example demonstrates how to cancel a running job.
|
||||
#
|
||||
# Requirements:
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash cancel.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI Job Cancellation Example"
|
||||
echo "============================================================"
|
||||
|
||||
# Submit a long-running job
|
||||
echo ""
|
||||
echo "1. Submitting a job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompt": "Write a detailed essay about the history of computing."}')
|
||||
|
||||
JOB_ID=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID"
|
||||
|
||||
# Wait briefly
|
||||
echo ""
|
||||
echo "2. Waiting 1 second..."
|
||||
sleep 1
|
||||
|
||||
# Cancel the job
|
||||
echo ""
|
||||
echo "3. Cancelling job..."
|
||||
CANCEL_RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs/$JOB_ID/cancel")
|
||||
STATUS=$(echo "$CANCEL_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['status'])")
|
||||
echo " Status after cancel: $STATUS"
|
||||
|
||||
# Verify cancellation
|
||||
echo ""
|
||||
echo "4. Verification:"
|
||||
if [ "$STATUS" = "cancelled" ]; then
|
||||
echo " ✓ SUCCESS: Job was cancelled"
|
||||
else
|
||||
echo " ℹ Job status: $STATUS (may have completed before cancel)"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Cancellation example completed!"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,75 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: Submit, Poll, and Get Result for Async Jobs
|
||||
#
|
||||
# This example demonstrates how to use the PraisonAI CLI to:
|
||||
# 1. Submit an async job
|
||||
# 2. Check job status
|
||||
# 3. Get the result
|
||||
#
|
||||
# Requirements:
|
||||
# - OPENAI_API_KEY environment variable set
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash cli_submit_poll_result.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI Async Jobs CLI Example"
|
||||
echo "============================================================"
|
||||
|
||||
# Check if server is running
|
||||
echo ""
|
||||
echo "Checking server health..."
|
||||
if ! curl -s "$API_URL/health" > /dev/null 2>&1; then
|
||||
echo "✗ Server not available at $API_URL"
|
||||
echo " Start the server with: python -m uvicorn praisonai.jobs.server:create_app --port 8005 --factory"
|
||||
exit 1
|
||||
fi
|
||||
echo "✓ Server is healthy"
|
||||
|
||||
# 1. Submit a job
|
||||
echo ""
|
||||
echo "1. Submitting job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompt": "What is 2+2? Answer with just the number."}')
|
||||
|
||||
JOB_ID=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID"
|
||||
|
||||
# 2. Poll for status
|
||||
echo ""
|
||||
echo "2. Polling for status..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/api/v1/runs/$JOB_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['status'])")
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['progress']['percentage'])")
|
||||
|
||||
echo " Status: $STATUS | Progress: ${PROGRESS}%"
|
||||
|
||||
if [ "$STATUS" = "succeeded" ] || [ "$STATUS" = "failed" ] || [ "$STATUS" = "cancelled" ]; then
|
||||
break
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
|
||||
# 3. Get result
|
||||
echo ""
|
||||
echo "3. Getting result..."
|
||||
if [ "$STATUS" = "succeeded" ]; then
|
||||
RESULT=$(curl -s "$API_URL/api/v1/runs/$JOB_ID/result")
|
||||
echo " Result: $(echo "$RESULT" | python3 -c "import sys,json; print(json.load(sys.stdin)['result'])")"
|
||||
echo " Duration: $(echo "$RESULT" | python3 -c "import sys,json; print(f\"{json.load(sys.stdin)['duration_seconds']:.2f}s\")")"
|
||||
else
|
||||
echo " Job did not succeed: $STATUS"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Example completed!"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,57 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: Idempotency Key Usage
|
||||
#
|
||||
# This example demonstrates how idempotency keys prevent duplicate job submissions.
|
||||
# Submitting the same request with the same idempotency key returns the existing job.
|
||||
#
|
||||
# Requirements:
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash idempotency.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
IDEM_KEY="example-idem-key-$(date +%s)"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI Idempotency Key Example"
|
||||
echo "============================================================"
|
||||
|
||||
# First submission
|
||||
echo ""
|
||||
echo "1. First submission with Idempotency-Key: $IDEM_KEY"
|
||||
RESPONSE1=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Idempotency-Key: $IDEM_KEY" \
|
||||
-d '{"prompt": "What is 1+1?"}')
|
||||
|
||||
JOB_ID1=$(echo "$RESPONSE1" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID1"
|
||||
|
||||
# Second submission with same key
|
||||
echo ""
|
||||
echo "2. Second submission with SAME Idempotency-Key: $IDEM_KEY"
|
||||
RESPONSE2=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Idempotency-Key: $IDEM_KEY" \
|
||||
-d '{"prompt": "Different prompt but same key"}')
|
||||
|
||||
JOB_ID2=$(echo "$RESPONSE2" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID2"
|
||||
|
||||
# Verify same job ID
|
||||
echo ""
|
||||
echo "3. Verification:"
|
||||
if [ "$JOB_ID1" = "$JOB_ID2" ]; then
|
||||
echo " ✓ SUCCESS: Same job ID returned (no duplicate created)"
|
||||
else
|
||||
echo " ✗ FAILURE: Different job IDs returned"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Idempotency example completed!"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,121 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SDK Example: Submit, Poll, and Get Result for Async Jobs
|
||||
|
||||
This example demonstrates how to use the PraisonAI Jobs SDK to:
|
||||
1. Submit an async job
|
||||
2. Poll for status
|
||||
3. Get the result
|
||||
|
||||
Requirements:
|
||||
- OPENAI_API_KEY environment variable set
|
||||
- Jobs server running on http://127.0.0.1:8005
|
||||
|
||||
Usage:
|
||||
python sdk_submit_poll_result.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import httpx
|
||||
|
||||
# Configuration
|
||||
API_URL = os.getenv("PRAISONAI_BASE_URL", "http://127.0.0.1:8005")
|
||||
POLL_INTERVAL = 2 # seconds
|
||||
|
||||
|
||||
def submit_job(prompt: str, idempotency_key: str = None) -> dict:
|
||||
"""Submit a job and return the response."""
|
||||
headers = {}
|
||||
if idempotency_key:
|
||||
headers["Idempotency-Key"] = idempotency_key
|
||||
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.post(
|
||||
f"{API_URL}/api/v1/runs",
|
||||
json={"prompt": prompt},
|
||||
headers=headers if headers else None
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_status(job_id: str) -> dict:
|
||||
"""Get job status."""
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.get(f"{API_URL}/api/v1/runs/{job_id}")
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_result(job_id: str) -> dict:
|
||||
"""Get job result."""
|
||||
with httpx.Client(timeout=30.0) as client:
|
||||
response = client.get(f"{API_URL}/api/v1/runs/{job_id}/result")
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def wait_for_completion(job_id: str, max_wait: int = 120) -> dict:
|
||||
"""Poll until job completes."""
|
||||
start = time.time()
|
||||
while time.time() - start < max_wait:
|
||||
status = get_status(job_id)
|
||||
print(f" Status: {status['status']} | Progress: {status['progress']['percentage']:.0f}%")
|
||||
|
||||
if status["status"] in ("succeeded", "failed", "cancelled"):
|
||||
return status
|
||||
|
||||
# Honor retry_after if present
|
||||
wait_time = status.get("retry_after") or POLL_INTERVAL
|
||||
time.sleep(wait_time)
|
||||
|
||||
raise TimeoutError(f"Job {job_id} did not complete within {max_wait}s")
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("PraisonAI Async Jobs SDK Example")
|
||||
print("=" * 60)
|
||||
|
||||
# Check if server is running
|
||||
try:
|
||||
with httpx.Client(timeout=5.0) as client:
|
||||
health = client.get(f"{API_URL}/health")
|
||||
health.raise_for_status()
|
||||
print(f"✓ Server healthy: {health.json()['status']}")
|
||||
except Exception as e:
|
||||
print(f"✗ Server not available at {API_URL}: {e}")
|
||||
print(" Start the server with: python -m uvicorn praisonai.jobs.server:create_app --port 8005 --factory")
|
||||
return
|
||||
|
||||
# 1. Submit a job
|
||||
print("\n1. Submitting job...")
|
||||
prompt = "What is the capital of France? Answer in one word."
|
||||
result = submit_job(prompt)
|
||||
job_id = result["job_id"]
|
||||
print(f" Job ID: {job_id}")
|
||||
print(f" Status: {result['status']}")
|
||||
print(f" Poll URL: {result['poll_url']}")
|
||||
|
||||
# 2. Wait for completion
|
||||
print("\n2. Waiting for completion...")
|
||||
final_status = wait_for_completion(job_id)
|
||||
print(f" Final status: {final_status['status']}")
|
||||
|
||||
# 3. Get result
|
||||
if final_status["status"] == "succeeded":
|
||||
print("\n3. Getting result...")
|
||||
result = get_result(job_id)
|
||||
print(f" Result: {result['result']}")
|
||||
print(f" Duration: {result['duration_seconds']:.2f}s")
|
||||
else:
|
||||
print(f"\n3. Job did not succeed: {final_status.get('error', 'Unknown error')}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Example completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,51 @@
|
||||
#!/bin/bash
|
||||
# CLI Example: SSE Streaming
|
||||
#
|
||||
# This example demonstrates how to stream job progress via Server-Sent Events.
|
||||
#
|
||||
# Requirements:
|
||||
# - Jobs server running on http://127.0.0.1:8005
|
||||
#
|
||||
# Usage:
|
||||
# bash sse_stream.sh
|
||||
|
||||
set -e
|
||||
|
||||
API_URL="${PRAISONAI_BASE_URL:-http://127.0.0.1:8005}"
|
||||
|
||||
echo "============================================================"
|
||||
echo "PraisonAI SSE Streaming Example"
|
||||
echo "============================================================"
|
||||
|
||||
# Submit a job
|
||||
echo ""
|
||||
echo "1. Submitting job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/api/v1/runs" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompt": "Count from 1 to 5."}')
|
||||
|
||||
JOB_ID=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
|
||||
echo " Job ID: $JOB_ID"
|
||||
|
||||
# Stream progress
|
||||
echo ""
|
||||
echo "2. Streaming progress (Ctrl+C to stop)..."
|
||||
echo " Connecting to: $API_URL/api/v1/runs/$JOB_ID/stream"
|
||||
echo ""
|
||||
|
||||
# Use curl to stream SSE events (timeout after 60s)
|
||||
curl -s -N --max-time 60 "$API_URL/api/v1/runs/$JOB_ID/stream" | while read -r line; do
|
||||
if [[ "$line" == data:* ]]; then
|
||||
data="${line#data: }"
|
||||
if [ "$data" = "[DONE]" ]; then
|
||||
echo " Stream completed"
|
||||
break
|
||||
fi
|
||||
echo " Event: $data"
|
||||
fi
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo "Streaming example completed!"
|
||||
echo "============================================================"
|
||||
@@ -43,8 +43,7 @@ from .tools.base import BaseTool, ToolResult, ToolValidationError, validate_tool
|
||||
from .tools.decorator import tool, FunctionTool
|
||||
from .tools.registry import get_registry, register_tool, get_tool, ToolRegistry
|
||||
from .agents.autoagents import AutoAgents
|
||||
from .knowledge.knowledge import Knowledge
|
||||
from .knowledge.chunking import Chunking
|
||||
# Knowledge is lazy-loaded to avoid importing heavy deps (chromadb, mem0) at startup
|
||||
# MCP support (optional)
|
||||
try:
|
||||
from .mcp.mcp import MCP
|
||||
@@ -53,7 +52,7 @@ except ImportError:
|
||||
_mcp_available = False
|
||||
MCP = None
|
||||
from .session import Session
|
||||
from .memory.memory import Memory
|
||||
# Memory is lazy-loaded to avoid importing chromadb at startup
|
||||
from .db import db
|
||||
from .obs import obs
|
||||
# Workflows - import from dedicated workflows module
|
||||
@@ -67,8 +66,16 @@ from .guardrails import GuardrailResult, LLMGuardrail
|
||||
|
||||
# Fast Context support (lazy loaded to avoid performance impact)
|
||||
def __getattr__(name):
|
||||
"""Lazy load FastContext to avoid impacting package load time."""
|
||||
if name == "FastContext":
|
||||
"""Lazy load heavy modules to avoid impacting package load time."""
|
||||
# Knowledge module (imports chromadb, mem0)
|
||||
if name == "Knowledge":
|
||||
from praisonaiagents.knowledge.knowledge import Knowledge
|
||||
return Knowledge
|
||||
elif name == "Chunking":
|
||||
from praisonaiagents.knowledge.chunking import Chunking
|
||||
return Chunking
|
||||
# FastContext support
|
||||
elif name == "FastContext":
|
||||
from praisonaiagents.context.fast import FastContext
|
||||
return FastContext
|
||||
elif name == "FastContextResult":
|
||||
@@ -93,6 +100,10 @@ def __getattr__(name):
|
||||
elif name == "SkillLoader":
|
||||
from praisonaiagents.skills import SkillLoader
|
||||
return SkillLoader
|
||||
# Memory module (imports chromadb)
|
||||
elif name == "Memory":
|
||||
from praisonaiagents.memory.memory import Memory
|
||||
return Memory
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
# Planning mode support
|
||||
|
||||
@@ -1,8 +1,130 @@
|
||||
"""
|
||||
PraisonAI Knowledge - Advanced knowledge management system with configurable features
|
||||
PraisonAI Knowledge - Advanced knowledge management system with configurable features.
|
||||
|
||||
This module provides:
|
||||
- Document readers (ReaderProtocol, ReaderRegistry)
|
||||
- Vector stores (VectorStoreProtocol, VectorStoreRegistry)
|
||||
- Retrieval strategies (RetrieverProtocol, RetrievalStrategy)
|
||||
- Rerankers (RerankerProtocol, RerankerRegistry)
|
||||
- Index types (IndexProtocol, IndexType)
|
||||
- Query engines (QueryEngineProtocol, QueryMode)
|
||||
"""
|
||||
|
||||
# Core Knowledge class (always available)
|
||||
from praisonaiagents.knowledge.knowledge import Knowledge
|
||||
from praisonaiagents.knowledge.chunking import Chunking
|
||||
|
||||
__all__ = ["Knowledge", "Chunking"]
|
||||
# Lazy loading for protocols and registries to avoid import overhead
|
||||
_LAZY_IMPORTS = {
|
||||
# Readers
|
||||
"Document": ("praisonaiagents.knowledge.readers", "Document"),
|
||||
"ReaderProtocol": ("praisonaiagents.knowledge.readers", "ReaderProtocol"),
|
||||
"ReaderRegistry": ("praisonaiagents.knowledge.readers", "ReaderRegistry"),
|
||||
"get_reader_registry": ("praisonaiagents.knowledge.readers", "get_reader_registry"),
|
||||
"detect_source_kind": ("praisonaiagents.knowledge.readers", "detect_source_kind"),
|
||||
|
||||
# Vector stores
|
||||
"VectorRecord": ("praisonaiagents.knowledge.vector_store", "VectorRecord"),
|
||||
"VectorStoreProtocol": ("praisonaiagents.knowledge.vector_store", "VectorStoreProtocol"),
|
||||
"VectorStoreRegistry": ("praisonaiagents.knowledge.vector_store", "VectorStoreRegistry"),
|
||||
"get_vector_store_registry": ("praisonaiagents.knowledge.vector_store", "get_vector_store_registry"),
|
||||
"InMemoryVectorStore": ("praisonaiagents.knowledge.vector_store", "InMemoryVectorStore"),
|
||||
|
||||
# Retrieval
|
||||
"RetrievalResult": ("praisonaiagents.knowledge.retrieval", "RetrievalResult"),
|
||||
"RetrievalStrategy": ("praisonaiagents.knowledge.retrieval", "RetrievalStrategy"),
|
||||
"RetrieverProtocol": ("praisonaiagents.knowledge.retrieval", "RetrieverProtocol"),
|
||||
"RetrieverRegistry": ("praisonaiagents.knowledge.retrieval", "RetrieverRegistry"),
|
||||
"get_retriever_registry": ("praisonaiagents.knowledge.retrieval", "get_retriever_registry"),
|
||||
"reciprocal_rank_fusion": ("praisonaiagents.knowledge.retrieval", "reciprocal_rank_fusion"),
|
||||
"merge_adjacent_chunks": ("praisonaiagents.knowledge.retrieval", "merge_adjacent_chunks"),
|
||||
|
||||
# Rerankers
|
||||
"RerankResult": ("praisonaiagents.knowledge.rerankers", "RerankResult"),
|
||||
"RerankerProtocol": ("praisonaiagents.knowledge.rerankers", "RerankerProtocol"),
|
||||
"RerankerRegistry": ("praisonaiagents.knowledge.rerankers", "RerankerRegistry"),
|
||||
"get_reranker_registry": ("praisonaiagents.knowledge.rerankers", "get_reranker_registry"),
|
||||
"SimpleReranker": ("praisonaiagents.knowledge.rerankers", "SimpleReranker"),
|
||||
|
||||
# Index
|
||||
"IndexType": ("praisonaiagents.knowledge.index", "IndexType"),
|
||||
"IndexStats": ("praisonaiagents.knowledge.index", "IndexStats"),
|
||||
"IndexProtocol": ("praisonaiagents.knowledge.index", "IndexProtocol"),
|
||||
"IndexRegistry": ("praisonaiagents.knowledge.index", "IndexRegistry"),
|
||||
"get_index_registry": ("praisonaiagents.knowledge.index", "get_index_registry"),
|
||||
"KeywordIndex": ("praisonaiagents.knowledge.index", "KeywordIndex"),
|
||||
|
||||
# Query engine
|
||||
"QueryMode": ("praisonaiagents.knowledge.query_engine", "QueryMode"),
|
||||
"QueryResult": ("praisonaiagents.knowledge.query_engine", "QueryResult"),
|
||||
"QueryEngineProtocol": ("praisonaiagents.knowledge.query_engine", "QueryEngineProtocol"),
|
||||
"QueryEngineRegistry": ("praisonaiagents.knowledge.query_engine", "QueryEngineRegistry"),
|
||||
"get_query_engine_registry": ("praisonaiagents.knowledge.query_engine", "get_query_engine_registry"),
|
||||
"decompose_question": ("praisonaiagents.knowledge.query_engine", "decompose_question"),
|
||||
"SimpleQueryEngine": ("praisonaiagents.knowledge.query_engine", "SimpleQueryEngine"),
|
||||
"SubQuestionEngine": ("praisonaiagents.knowledge.query_engine", "SubQuestionEngine"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load protocols and registries."""
|
||||
if name in _LAZY_IMPORTS:
|
||||
module_path, attr_name = _LAZY_IMPORTS[name]
|
||||
import importlib
|
||||
module = importlib.import_module(module_path)
|
||||
return getattr(module, attr_name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__():
|
||||
"""List available attributes."""
|
||||
return list(_LAZY_IMPORTS.keys()) + ["Knowledge", "Chunking"]
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Core
|
||||
"Knowledge",
|
||||
"Chunking",
|
||||
# Readers
|
||||
"Document",
|
||||
"ReaderProtocol",
|
||||
"ReaderRegistry",
|
||||
"get_reader_registry",
|
||||
"detect_source_kind",
|
||||
# Vector stores
|
||||
"VectorRecord",
|
||||
"VectorStoreProtocol",
|
||||
"VectorStoreRegistry",
|
||||
"get_vector_store_registry",
|
||||
"InMemoryVectorStore",
|
||||
# Retrieval
|
||||
"RetrievalResult",
|
||||
"RetrievalStrategy",
|
||||
"RetrieverProtocol",
|
||||
"RetrieverRegistry",
|
||||
"get_retriever_registry",
|
||||
"reciprocal_rank_fusion",
|
||||
"merge_adjacent_chunks",
|
||||
# Rerankers
|
||||
"RerankResult",
|
||||
"RerankerProtocol",
|
||||
"RerankerRegistry",
|
||||
"get_reranker_registry",
|
||||
"SimpleReranker",
|
||||
# Index
|
||||
"IndexType",
|
||||
"IndexStats",
|
||||
"IndexProtocol",
|
||||
"IndexRegistry",
|
||||
"get_index_registry",
|
||||
"KeywordIndex",
|
||||
# Query engine
|
||||
"QueryMode",
|
||||
"QueryResult",
|
||||
"QueryEngineProtocol",
|
||||
"QueryEngineRegistry",
|
||||
"get_query_engine_registry",
|
||||
"decompose_question",
|
||||
"SimpleQueryEngine",
|
||||
"SubQuestionEngine",
|
||||
]
|
||||
@@ -0,0 +1,332 @@
|
||||
"""
|
||||
Index Type Selection for PraisonAI Agents.
|
||||
|
||||
This module defines index types and protocols.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from enum import Enum
|
||||
import logging
|
||||
import math
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IndexType(str, Enum):
|
||||
"""Available index types."""
|
||||
VECTOR = "vector"
|
||||
KEYWORD = "keyword"
|
||||
HYBRID = "hybrid"
|
||||
GRAPH = "graph"
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexStats:
|
||||
"""Statistics about an index."""
|
||||
index_type: IndexType
|
||||
document_count: int
|
||||
chunk_count: int
|
||||
total_tokens: int = 0
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class IndexProtocol(Protocol):
|
||||
"""
|
||||
Protocol for document indices.
|
||||
|
||||
Implementations must provide methods for indexing and querying.
|
||||
"""
|
||||
|
||||
name: str
|
||||
index_type: IndexType
|
||||
|
||||
def add_documents(
|
||||
self,
|
||||
texts: List[str],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
embeddings: Optional[List[List[float]]] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
Add documents to the index.
|
||||
|
||||
Args:
|
||||
texts: List of text content
|
||||
metadatas: Optional list of metadata dicts
|
||||
ids: Optional list of IDs
|
||||
embeddings: Optional pre-computed embeddings (for vector index)
|
||||
|
||||
Returns:
|
||||
List of document IDs
|
||||
"""
|
||||
...
|
||||
|
||||
def query(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
query_embedding: Optional[List[float]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query the index.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
top_k: Number of results
|
||||
filter: Optional metadata filter
|
||||
query_embedding: Optional pre-computed query embedding
|
||||
|
||||
Returns:
|
||||
List of results with text, score, metadata
|
||||
"""
|
||||
...
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete documents from the index."""
|
||||
...
|
||||
|
||||
def stats(self) -> IndexStats:
|
||||
"""Get index statistics."""
|
||||
...
|
||||
|
||||
|
||||
class IndexRegistry:
|
||||
"""Registry for index types."""
|
||||
|
||||
_instance: Optional["IndexRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "IndexRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._indices: Dict[str, Callable[..., IndexProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(self, name: str, factory: Callable[..., IndexProtocol]) -> None:
|
||||
"""Register an index factory."""
|
||||
self._indices[name] = factory
|
||||
|
||||
def get(self, name: str, **kwargs) -> Optional[IndexProtocol]:
|
||||
"""Get an index by name."""
|
||||
if name in self._indices:
|
||||
try:
|
||||
return self._indices[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize index '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_indices(self) -> List[str]:
|
||||
"""List all registered index names."""
|
||||
return list(self._indices.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered indices."""
|
||||
self._indices.clear()
|
||||
|
||||
|
||||
def get_index_registry() -> IndexRegistry:
|
||||
"""Get the global index registry instance."""
|
||||
return IndexRegistry()
|
||||
|
||||
|
||||
class KeywordIndex:
|
||||
"""
|
||||
Simple BM25-like keyword index.
|
||||
|
||||
Pure Python implementation with no external dependencies.
|
||||
"""
|
||||
|
||||
name: str = "keyword"
|
||||
index_type: IndexType = IndexType.KEYWORD
|
||||
|
||||
def __init__(self, k1: float = 1.5, b: float = 0.75, **kwargs):
|
||||
self.k1 = k1
|
||||
self.b = b
|
||||
self._documents: Dict[str, Dict[str, Any]] = {} # id -> {text, metadata, terms}
|
||||
self._term_doc_freq: Dict[str, int] = {} # term -> doc count
|
||||
self._avg_doc_len: float = 0.0
|
||||
self._total_docs: int = 0
|
||||
|
||||
def add_documents(
|
||||
self,
|
||||
texts: List[str],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
embeddings: Optional[List[List[float]]] = None
|
||||
) -> List[str]:
|
||||
"""Add documents to the keyword index."""
|
||||
import uuid
|
||||
|
||||
metadatas = metadatas or [{} for _ in texts]
|
||||
ids = ids or [str(uuid.uuid4()) for _ in texts]
|
||||
|
||||
total_len = sum(len(doc["text"].split()) for doc in self._documents.values())
|
||||
|
||||
for text, metadata, doc_id in zip(texts, metadatas, ids):
|
||||
terms = self._tokenize(text)
|
||||
term_freq = {}
|
||||
for term in terms:
|
||||
term_freq[term] = term_freq.get(term, 0) + 1
|
||||
|
||||
self._documents[doc_id] = {
|
||||
"text": text,
|
||||
"metadata": metadata,
|
||||
"terms": term_freq,
|
||||
"length": len(terms)
|
||||
}
|
||||
|
||||
# Update document frequency
|
||||
for term in set(terms):
|
||||
self._term_doc_freq[term] = self._term_doc_freq.get(term, 0) + 1
|
||||
|
||||
total_len += len(terms)
|
||||
|
||||
self._total_docs = len(self._documents)
|
||||
self._avg_doc_len = total_len / self._total_docs if self._total_docs > 0 else 0
|
||||
|
||||
return ids
|
||||
|
||||
def query(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
query_embedding: Optional[List[float]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Query using BM25 scoring."""
|
||||
query_terms = self._tokenize(query)
|
||||
|
||||
scores = []
|
||||
for doc_id, doc in self._documents.items():
|
||||
# Apply filter
|
||||
if filter:
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if doc["metadata"].get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if not match:
|
||||
continue
|
||||
|
||||
score = self._bm25_score(query_terms, doc)
|
||||
scores.append({
|
||||
"id": doc_id,
|
||||
"text": doc["text"],
|
||||
"score": score,
|
||||
"metadata": doc["metadata"]
|
||||
})
|
||||
|
||||
# Sort by score
|
||||
scores.sort(key=lambda x: x["score"], reverse=True)
|
||||
return scores[:top_k]
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete documents from the index."""
|
||||
if delete_all:
|
||||
count = len(self._documents)
|
||||
self._documents.clear()
|
||||
self._term_doc_freq.clear()
|
||||
self._total_docs = 0
|
||||
self._avg_doc_len = 0
|
||||
return count
|
||||
|
||||
deleted = 0
|
||||
to_delete = []
|
||||
|
||||
if ids:
|
||||
to_delete = [id_ for id_ in ids if id_ in self._documents]
|
||||
elif filter:
|
||||
for doc_id, doc in self._documents.items():
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if doc["metadata"].get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if match:
|
||||
to_delete.append(doc_id)
|
||||
|
||||
for doc_id in to_delete:
|
||||
doc = self._documents[doc_id]
|
||||
# Update term frequencies
|
||||
for term in doc["terms"]:
|
||||
if term in self._term_doc_freq:
|
||||
self._term_doc_freq[term] -= 1
|
||||
if self._term_doc_freq[term] <= 0:
|
||||
del self._term_doc_freq[term]
|
||||
del self._documents[doc_id]
|
||||
deleted += 1
|
||||
|
||||
self._total_docs = len(self._documents)
|
||||
if self._total_docs > 0:
|
||||
self._avg_doc_len = sum(d["length"] for d in self._documents.values()) / self._total_docs
|
||||
else:
|
||||
self._avg_doc_len = 0
|
||||
|
||||
return deleted
|
||||
|
||||
def stats(self) -> IndexStats:
|
||||
"""Get index statistics."""
|
||||
return IndexStats(
|
||||
index_type=self.index_type,
|
||||
document_count=self._total_docs,
|
||||
chunk_count=self._total_docs,
|
||||
total_tokens=sum(d["length"] for d in self._documents.values()),
|
||||
metadata={"unique_terms": len(self._term_doc_freq)}
|
||||
)
|
||||
|
||||
def _tokenize(self, text: str) -> List[str]:
|
||||
"""Simple tokenization."""
|
||||
import re
|
||||
# Lowercase and split on non-alphanumeric
|
||||
tokens = re.findall(r'\b\w+\b', text.lower())
|
||||
# Remove very short tokens
|
||||
return [t for t in tokens if len(t) > 1]
|
||||
|
||||
def _bm25_score(self, query_terms: List[str], doc: Dict[str, Any]) -> float:
|
||||
"""Calculate BM25 score."""
|
||||
score = 0.0
|
||||
doc_len = doc["length"]
|
||||
term_freq = doc["terms"]
|
||||
|
||||
for term in query_terms:
|
||||
if term not in term_freq:
|
||||
continue
|
||||
|
||||
tf = term_freq[term]
|
||||
df = self._term_doc_freq.get(term, 0)
|
||||
|
||||
# IDF component
|
||||
idf = math.log((self._total_docs - df + 0.5) / (df + 0.5) + 1)
|
||||
|
||||
# TF component with length normalization
|
||||
tf_norm = (tf * (self.k1 + 1)) / (tf + self.k1 * (1 - self.b + self.b * doc_len / self._avg_doc_len))
|
||||
|
||||
score += idf * tf_norm
|
||||
|
||||
return score
|
||||
|
||||
|
||||
# Register the keyword index by default
|
||||
def _register_default_indices():
|
||||
"""Register default indices."""
|
||||
registry = get_index_registry()
|
||||
registry.register("keyword", KeywordIndex)
|
||||
|
||||
|
||||
_register_default_indices()
|
||||
@@ -0,0 +1,351 @@
|
||||
"""
|
||||
Query Engine Patterns for PraisonAI Agents.
|
||||
|
||||
This module defines query modes and protocols.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from enum import Enum
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class QueryMode(str, Enum):
|
||||
"""Available query modes."""
|
||||
DEFAULT = "default"
|
||||
SUB_QUESTION = "sub_question"
|
||||
SQL = "sql"
|
||||
ROUTER = "router"
|
||||
SUMMARIZE = "summarize"
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
"""
|
||||
Result from a query operation.
|
||||
|
||||
Attributes:
|
||||
answer: The synthesized answer
|
||||
sources: List of source documents/chunks used
|
||||
sub_questions: Optional list of sub-questions (for sub_question mode)
|
||||
metadata: Additional metadata
|
||||
"""
|
||||
answer: str
|
||||
sources: List[Dict[str, Any]] = field(default_factory=list)
|
||||
sub_questions: Optional[List[str]] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"answer": self.answer,
|
||||
"sources": self.sources,
|
||||
"sub_questions": self.sub_questions,
|
||||
"metadata": self.metadata
|
||||
}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class QueryEngineProtocol(Protocol):
|
||||
"""
|
||||
Protocol for query engines.
|
||||
|
||||
Implementations must provide a query method that takes a question
|
||||
and returns an answer with sources.
|
||||
"""
|
||||
|
||||
name: str
|
||||
mode: QueryMode
|
||||
|
||||
def query(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""
|
||||
Query the knowledge base.
|
||||
|
||||
Args:
|
||||
question: The user's question
|
||||
context: Optional pre-retrieved context
|
||||
**kwargs: Additional engine-specific options
|
||||
|
||||
Returns:
|
||||
QueryResult with answer and sources
|
||||
"""
|
||||
...
|
||||
|
||||
async def aquery(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Async version of query."""
|
||||
...
|
||||
|
||||
|
||||
class QueryEngineRegistry:
|
||||
"""Registry for query engines."""
|
||||
|
||||
_instance: Optional["QueryEngineRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "QueryEngineRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._engines: Dict[str, Callable[..., QueryEngineProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(self, name: str, factory: Callable[..., QueryEngineProtocol]) -> None:
|
||||
"""Register a query engine factory."""
|
||||
self._engines[name] = factory
|
||||
|
||||
def get(self, name: str, **kwargs) -> Optional[QueryEngineProtocol]:
|
||||
"""Get a query engine by name."""
|
||||
if name in self._engines:
|
||||
try:
|
||||
return self._engines[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize query engine '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_engines(self) -> List[str]:
|
||||
"""List all registered engine names."""
|
||||
return list(self._engines.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered engines."""
|
||||
self._engines.clear()
|
||||
|
||||
|
||||
def get_query_engine_registry() -> QueryEngineRegistry:
|
||||
"""Get the global query engine registry instance."""
|
||||
return QueryEngineRegistry()
|
||||
|
||||
|
||||
def decompose_question(question: str) -> List[str]:
|
||||
"""
|
||||
Simple deterministic question decomposition.
|
||||
|
||||
Splits complex questions into sub-questions without LLM.
|
||||
Used as fallback when LLM is not available.
|
||||
|
||||
Args:
|
||||
question: The complex question
|
||||
|
||||
Returns:
|
||||
List of sub-questions
|
||||
"""
|
||||
sub_questions = []
|
||||
|
||||
# Split on common conjunctions
|
||||
conjunctions = [" and ", " also ", " additionally ", " furthermore ", " moreover "]
|
||||
parts = [question]
|
||||
|
||||
for conj in conjunctions:
|
||||
new_parts = []
|
||||
for part in parts:
|
||||
new_parts.extend(part.split(conj))
|
||||
parts = new_parts
|
||||
|
||||
# Clean up parts
|
||||
for part in parts:
|
||||
part = part.strip()
|
||||
if part:
|
||||
# Ensure it ends with a question mark
|
||||
if not part.endswith("?"):
|
||||
part = part + "?"
|
||||
# Capitalize first letter
|
||||
part = part[0].upper() + part[1:] if len(part) > 1 else part.upper()
|
||||
sub_questions.append(part)
|
||||
|
||||
# If no decomposition happened, return original
|
||||
if len(sub_questions) <= 1:
|
||||
return [question]
|
||||
|
||||
return sub_questions
|
||||
|
||||
|
||||
def synthesize_answer(
|
||||
question: str,
|
||||
contexts: List[str],
|
||||
max_context_length: int = 4000
|
||||
) -> str:
|
||||
"""
|
||||
Simple answer synthesis without LLM.
|
||||
|
||||
Concatenates relevant context with the question.
|
||||
Used as fallback when LLM is not available.
|
||||
|
||||
Args:
|
||||
question: The original question
|
||||
contexts: List of context strings
|
||||
max_context_length: Maximum total context length
|
||||
|
||||
Returns:
|
||||
Synthesized answer string
|
||||
"""
|
||||
if not contexts:
|
||||
return "No relevant information found."
|
||||
|
||||
# Truncate contexts to fit
|
||||
total_length = 0
|
||||
selected_contexts = []
|
||||
|
||||
for ctx in contexts:
|
||||
if total_length + len(ctx) > max_context_length:
|
||||
remaining = max_context_length - total_length
|
||||
if remaining > 100:
|
||||
selected_contexts.append(ctx[:remaining] + "...")
|
||||
break
|
||||
selected_contexts.append(ctx)
|
||||
total_length += len(ctx)
|
||||
|
||||
# Format as a simple answer
|
||||
context_text = "\n\n".join(selected_contexts)
|
||||
|
||||
return f"Based on the available information:\n\n{context_text}"
|
||||
|
||||
|
||||
class SimpleQueryEngine:
|
||||
"""
|
||||
Simple query engine for testing and fallback.
|
||||
|
||||
Uses basic context concatenation without LLM synthesis.
|
||||
"""
|
||||
|
||||
name: str = "simple"
|
||||
mode: QueryMode = QueryMode.DEFAULT
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def query(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Simple query that returns context as answer."""
|
||||
contexts = context or []
|
||||
answer = synthesize_answer(question, contexts)
|
||||
|
||||
return QueryResult(
|
||||
answer=answer,
|
||||
sources=[{"text": ctx} for ctx in contexts],
|
||||
metadata={"mode": self.mode.value}
|
||||
)
|
||||
|
||||
async def aquery(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.query(question, context, **kwargs)
|
||||
|
||||
|
||||
class SubQuestionEngine:
|
||||
"""
|
||||
Sub-question decomposition query engine.
|
||||
|
||||
Decomposes complex questions into simpler sub-questions.
|
||||
Uses deterministic fallback when LLM is not available.
|
||||
"""
|
||||
|
||||
name: str = "sub_question"
|
||||
mode: QueryMode = QueryMode.SUB_QUESTION
|
||||
|
||||
def __init__(self, llm: Optional[Any] = None, **kwargs):
|
||||
self.llm = llm
|
||||
|
||||
def query(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
retriever: Optional[Any] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Query with sub-question decomposition."""
|
||||
# Decompose question
|
||||
if self.llm:
|
||||
sub_questions = self._decompose_with_llm(question)
|
||||
else:
|
||||
sub_questions = decompose_question(question)
|
||||
|
||||
# Gather context for each sub-question
|
||||
all_contexts = list(context or [])
|
||||
sub_answers = []
|
||||
|
||||
for sub_q in sub_questions:
|
||||
if retriever:
|
||||
# Retrieve for each sub-question
|
||||
try:
|
||||
results = retriever.retrieve(sub_q, top_k=3)
|
||||
for r in results:
|
||||
if hasattr(r, 'text'):
|
||||
all_contexts.append(r.text)
|
||||
elif isinstance(r, dict):
|
||||
all_contexts.append(r.get('text', str(r)))
|
||||
except Exception as e:
|
||||
logger.warning(f"Retrieval failed for sub-question: {e}")
|
||||
|
||||
# Simple sub-answer (would use LLM in full implementation)
|
||||
sub_answers.append(f"Sub-question: {sub_q}")
|
||||
|
||||
# Synthesize final answer
|
||||
answer = synthesize_answer(question, all_contexts)
|
||||
|
||||
return QueryResult(
|
||||
answer=answer,
|
||||
sources=[{"text": ctx} for ctx in all_contexts[:10]],
|
||||
sub_questions=sub_questions,
|
||||
metadata={"mode": self.mode.value, "sub_question_count": len(sub_questions)}
|
||||
)
|
||||
|
||||
async def aquery(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[List[str]] = None,
|
||||
**kwargs
|
||||
) -> QueryResult:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.query(question, context, **kwargs)
|
||||
|
||||
def _decompose_with_llm(self, question: str) -> List[str]:
|
||||
"""Decompose question using LLM."""
|
||||
try:
|
||||
prompt = f"""Break down this complex question into simpler sub-questions that can be answered independently.
|
||||
|
||||
Question: {question}
|
||||
|
||||
Return only the sub-questions, one per line, without numbering or bullets."""
|
||||
|
||||
response = self.llm.chat(prompt)
|
||||
if response:
|
||||
lines = response.strip().split("\n")
|
||||
return [line.strip() for line in lines if line.strip()]
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM decomposition failed: {e}")
|
||||
|
||||
return decompose_question(question)
|
||||
|
||||
|
||||
# Register default engines
|
||||
def _register_default_engines():
|
||||
"""Register default query engines."""
|
||||
registry = get_query_engine_registry()
|
||||
registry.register("default", SimpleQueryEngine)
|
||||
registry.register("simple", SimpleQueryEngine)
|
||||
registry.register("sub_question", SubQuestionEngine)
|
||||
|
||||
|
||||
_register_default_engines()
|
||||
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
Data Reader Adapter Protocol for PraisonAI Agents.
|
||||
|
||||
This module defines the protocol for document readers and a registry system.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer (praisonai.adapters.readers).
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
import os
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""
|
||||
Lightweight document representation.
|
||||
|
||||
Attributes:
|
||||
content: The text content of the document
|
||||
metadata: Optional metadata dict (source, filename, page, etc.)
|
||||
doc_id: Optional unique identifier
|
||||
embedding: Optional pre-computed embedding
|
||||
"""
|
||||
content: str
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
doc_id: Optional[str] = None
|
||||
embedding: Optional[List[float]] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.doc_id is None:
|
||||
import uuid
|
||||
self.doc_id = str(uuid.uuid4())
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"content": self.content,
|
||||
"metadata": self.metadata,
|
||||
"doc_id": self.doc_id,
|
||||
"embedding": self.embedding
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Document":
|
||||
"""Create from dictionary."""
|
||||
return cls(
|
||||
content=data.get("content", ""),
|
||||
metadata=data.get("metadata", {}),
|
||||
doc_id=data.get("doc_id"),
|
||||
embedding=data.get("embedding")
|
||||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ReaderProtocol(Protocol):
|
||||
"""
|
||||
Protocol for document readers.
|
||||
|
||||
Implementations must provide:
|
||||
- name: Identifier for the reader
|
||||
- supported_extensions: List of file extensions this reader handles
|
||||
- load(): Method to load documents from a source
|
||||
"""
|
||||
|
||||
name: str
|
||||
supported_extensions: List[str]
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Document]:
|
||||
"""
|
||||
Load documents from a source.
|
||||
|
||||
Args:
|
||||
source: File path, URL, or other source identifier
|
||||
metadata: Optional metadata to attach to all documents
|
||||
|
||||
Returns:
|
||||
List of Document objects
|
||||
"""
|
||||
...
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""
|
||||
Check if this reader can handle the given source.
|
||||
|
||||
Args:
|
||||
source: File path, URL, or other source identifier
|
||||
|
||||
Returns:
|
||||
True if this reader can handle the source
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class ReaderRegistry:
|
||||
"""
|
||||
Registry for document readers.
|
||||
|
||||
Provides lazy loading and automatic reader selection based on source type.
|
||||
"""
|
||||
|
||||
_instance: Optional["ReaderRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "ReaderRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._readers: Dict[str, Callable[[], ReaderProtocol]] = {}
|
||||
cls._instance._extension_map: Dict[str, str] = {}
|
||||
cls._instance._initialized_readers: Dict[str, ReaderProtocol] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[[], ReaderProtocol],
|
||||
extensions: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""
|
||||
Register a reader factory.
|
||||
|
||||
Args:
|
||||
name: Unique name for the reader
|
||||
factory: Callable that returns a ReaderProtocol instance
|
||||
extensions: List of file extensions this reader handles
|
||||
"""
|
||||
self._readers[name] = factory
|
||||
if extensions:
|
||||
for ext in extensions:
|
||||
ext_lower = ext.lower().lstrip(".")
|
||||
self._extension_map[ext_lower] = name
|
||||
|
||||
def get(self, name: str) -> Optional[ReaderProtocol]:
|
||||
"""
|
||||
Get a reader by name (lazy initialization).
|
||||
|
||||
Args:
|
||||
name: Reader name
|
||||
|
||||
Returns:
|
||||
ReaderProtocol instance or None
|
||||
"""
|
||||
if name in self._initialized_readers:
|
||||
return self._initialized_readers[name]
|
||||
|
||||
if name in self._readers:
|
||||
try:
|
||||
reader = self._readers[name]()
|
||||
self._initialized_readers[name] = reader
|
||||
return reader
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize reader '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def get_for_source(self, source: str) -> Optional[ReaderProtocol]:
|
||||
"""
|
||||
Get the best reader for a given source.
|
||||
|
||||
Args:
|
||||
source: File path, URL, or other source identifier
|
||||
|
||||
Returns:
|
||||
Best matching ReaderProtocol or None
|
||||
"""
|
||||
source_kind = detect_source_kind(source)
|
||||
|
||||
if source_kind == "file":
|
||||
ext = os.path.splitext(source)[1].lower().lstrip(".")
|
||||
if ext in self._extension_map:
|
||||
return self.get(self._extension_map[ext])
|
||||
elif source_kind == "url":
|
||||
# Try URL reader
|
||||
if "url" in self._readers:
|
||||
return self.get("url")
|
||||
elif source_kind == "directory":
|
||||
if "directory" in self._readers:
|
||||
return self.get("directory")
|
||||
|
||||
# Fallback to auto reader if available
|
||||
if "auto" in self._readers:
|
||||
return self.get("auto")
|
||||
|
||||
return None
|
||||
|
||||
def list_readers(self) -> List[str]:
|
||||
"""List all registered reader names."""
|
||||
return list(self._readers.keys())
|
||||
|
||||
def list_extensions(self) -> Dict[str, str]:
|
||||
"""List all registered extensions and their readers."""
|
||||
return dict(self._extension_map)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered readers (mainly for testing)."""
|
||||
self._readers.clear()
|
||||
self._extension_map.clear()
|
||||
self._initialized_readers.clear()
|
||||
|
||||
|
||||
def get_reader_registry() -> ReaderRegistry:
|
||||
"""Get the global reader registry instance."""
|
||||
return ReaderRegistry()
|
||||
|
||||
|
||||
def detect_source_kind(source: str) -> str:
|
||||
"""
|
||||
Detect the kind of source without importing heavy libraries.
|
||||
|
||||
Args:
|
||||
source: File path, URL, directory, or glob pattern
|
||||
|
||||
Returns:
|
||||
One of: "file", "url", "directory", "glob", "unknown"
|
||||
"""
|
||||
if not isinstance(source, str):
|
||||
return "unknown"
|
||||
|
||||
source = source.strip()
|
||||
|
||||
# Check for URL
|
||||
if source.startswith(("http://", "https://", "ftp://", "s3://", "gs://")):
|
||||
return "url"
|
||||
|
||||
# Check for glob pattern
|
||||
if "*" in source or "?" in source or "[" in source:
|
||||
return "glob"
|
||||
|
||||
# Check if it's a directory
|
||||
if os.path.isdir(source):
|
||||
return "directory"
|
||||
|
||||
# Check if it's a file
|
||||
if os.path.isfile(source):
|
||||
return "file"
|
||||
|
||||
# Check if it looks like a file path (has extension)
|
||||
if os.path.splitext(source)[1]:
|
||||
return "file"
|
||||
|
||||
return "unknown"
|
||||
|
||||
|
||||
def get_file_extension(source: str) -> str:
|
||||
"""
|
||||
Get the file extension from a source path.
|
||||
|
||||
Args:
|
||||
source: File path or URL
|
||||
|
||||
Returns:
|
||||
Lowercase extension without dot, or empty string
|
||||
"""
|
||||
# Handle URLs by extracting path
|
||||
if source.startswith(("http://", "https://")):
|
||||
from urllib.parse import urlparse
|
||||
path = urlparse(source).path
|
||||
return os.path.splitext(path)[1].lower().lstrip(".")
|
||||
|
||||
return os.path.splitext(source)[1].lower().lstrip(".")
|
||||
|
||||
|
||||
# Default extension mappings (reader name -> extensions)
|
||||
DEFAULT_EXTENSION_MAPPINGS = {
|
||||
"text": ["txt", "text"],
|
||||
"markdown": ["md", "markdown"],
|
||||
"json": ["json", "jsonl"],
|
||||
"csv": ["csv", "tsv"],
|
||||
"html": ["html", "htm"],
|
||||
"xml": ["xml"],
|
||||
"pdf": ["pdf"],
|
||||
"docx": ["docx", "doc"],
|
||||
"xlsx": ["xlsx", "xls"],
|
||||
"pptx": ["pptx", "ppt"],
|
||||
"image": ["jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp"],
|
||||
"audio": ["mp3", "wav", "ogg", "m4a", "flac"],
|
||||
"video": ["mp4", "avi", "mov", "mkv", "webm"],
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
Reranker Protocol for PraisonAI Agents.
|
||||
|
||||
This module defines the protocol for rerankers and a registry system.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by praisonai-tools or wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankResult:
|
||||
"""
|
||||
Result from a reranking operation.
|
||||
|
||||
Attributes:
|
||||
text: The text content
|
||||
score: Reranking score (higher is more relevant)
|
||||
original_index: Original position before reranking
|
||||
metadata: Optional metadata
|
||||
"""
|
||||
text: str
|
||||
score: float
|
||||
original_index: int
|
||||
metadata: Dict[str, Any] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.metadata is None:
|
||||
self.metadata = {}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RerankerProtocol(Protocol):
|
||||
"""
|
||||
Protocol for rerankers.
|
||||
|
||||
Implementations must provide a rerank method that takes a query
|
||||
and a list of documents and returns reranked results.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""
|
||||
Rerank documents based on relevance to query.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of document texts to rerank
|
||||
top_k: Optional limit on results (None = return all)
|
||||
**kwargs: Additional reranker-specific options
|
||||
|
||||
Returns:
|
||||
List of RerankResult objects sorted by relevance
|
||||
"""
|
||||
...
|
||||
|
||||
async def arerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""Async version of rerank."""
|
||||
...
|
||||
|
||||
|
||||
class RerankerRegistry:
|
||||
"""
|
||||
Registry for rerankers.
|
||||
"""
|
||||
|
||||
_instance: Optional["RerankerRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "RerankerRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._rerankers: Dict[str, Callable[..., RerankerProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[..., RerankerProtocol]
|
||||
) -> None:
|
||||
"""Register a reranker factory."""
|
||||
self._rerankers[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
**kwargs
|
||||
) -> Optional[RerankerProtocol]:
|
||||
"""Get a reranker by name."""
|
||||
if name in self._rerankers:
|
||||
try:
|
||||
return self._rerankers[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize reranker '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_rerankers(self) -> List[str]:
|
||||
"""List all registered reranker names."""
|
||||
return list(self._rerankers.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered rerankers."""
|
||||
self._rerankers.clear()
|
||||
|
||||
|
||||
def get_reranker_registry() -> RerankerRegistry:
|
||||
"""Get the global reranker registry instance."""
|
||||
return RerankerRegistry()
|
||||
|
||||
|
||||
class SimpleReranker:
|
||||
"""
|
||||
Simple keyword-based reranker for testing and fallback.
|
||||
|
||||
Uses term frequency overlap for scoring. No external dependencies.
|
||||
"""
|
||||
|
||||
name: str = "simple"
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""Rerank using simple term overlap."""
|
||||
query_terms = set(query.lower().split())
|
||||
|
||||
results = []
|
||||
for i, doc in enumerate(documents):
|
||||
doc_terms = set(doc.lower().split())
|
||||
|
||||
# Calculate Jaccard similarity
|
||||
if not query_terms or not doc_terms:
|
||||
score = 0.0
|
||||
else:
|
||||
intersection = len(query_terms & doc_terms)
|
||||
union = len(query_terms | doc_terms)
|
||||
score = intersection / union if union > 0 else 0.0
|
||||
|
||||
results.append(RerankResult(
|
||||
text=doc,
|
||||
score=score,
|
||||
original_index=i
|
||||
))
|
||||
|
||||
# Sort by score descending
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
|
||||
if top_k is not None:
|
||||
results = results[:top_k]
|
||||
|
||||
return results
|
||||
|
||||
async def arerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[RerankResult]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.rerank(query, documents, top_k, **kwargs)
|
||||
|
||||
|
||||
# Register the simple reranker by default
|
||||
def _register_default_rerankers():
|
||||
"""Register default rerankers."""
|
||||
registry = get_reranker_registry()
|
||||
registry.register("simple", SimpleReranker)
|
||||
|
||||
|
||||
_register_default_rerankers()
|
||||
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
Retrieval Strategy Patterns for PraisonAI Agents.
|
||||
|
||||
This module defines protocols and strategies for document retrieval.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by the wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from enum import Enum
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RetrievalStrategy(str, Enum):
|
||||
"""Available retrieval strategies."""
|
||||
BASIC = "basic"
|
||||
FUSION = "fusion"
|
||||
RECURSIVE = "recursive"
|
||||
AUTO_MERGE = "auto_merge"
|
||||
HYBRID = "hybrid"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetrievalResult:
|
||||
"""
|
||||
Result from a retrieval operation.
|
||||
|
||||
Attributes:
|
||||
text: The retrieved text content
|
||||
score: Relevance score (0-1)
|
||||
metadata: Source metadata
|
||||
doc_id: Document identifier
|
||||
chunk_index: Index of chunk within document
|
||||
"""
|
||||
text: str
|
||||
score: float
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
doc_id: Optional[str] = None
|
||||
chunk_index: Optional[int] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"text": self.text,
|
||||
"score": self.score,
|
||||
"metadata": self.metadata,
|
||||
"doc_id": self.doc_id,
|
||||
"chunk_index": self.chunk_index
|
||||
}
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RetrieverProtocol(Protocol):
|
||||
"""
|
||||
Protocol for document retrievers.
|
||||
|
||||
Implementations must provide a retrieve method that takes a query
|
||||
and returns relevant documents.
|
||||
"""
|
||||
|
||||
name: str
|
||||
strategy: RetrievalStrategy
|
||||
|
||||
def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
Retrieve relevant documents for a query.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
top_k: Maximum number of results
|
||||
filter: Optional metadata filter
|
||||
**kwargs: Additional retriever-specific options
|
||||
|
||||
Returns:
|
||||
List of RetrievalResult objects sorted by relevance
|
||||
"""
|
||||
...
|
||||
|
||||
async def aretrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = 10,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[RetrievalResult]:
|
||||
"""Async version of retrieve."""
|
||||
...
|
||||
|
||||
|
||||
class RetrieverRegistry:
|
||||
"""
|
||||
Registry for retrieval strategies.
|
||||
"""
|
||||
|
||||
_instance: Optional["RetrieverRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "RetrieverRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._retrievers: Dict[str, Callable[..., RetrieverProtocol]] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[..., RetrieverProtocol]
|
||||
) -> None:
|
||||
"""Register a retriever factory."""
|
||||
self._retrievers[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
**kwargs
|
||||
) -> Optional[RetrieverProtocol]:
|
||||
"""Get a retriever by name."""
|
||||
if name in self._retrievers:
|
||||
try:
|
||||
return self._retrievers[name](**kwargs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize retriever '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_retrievers(self) -> List[str]:
|
||||
"""List all registered retriever names."""
|
||||
return list(self._retrievers.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered retrievers."""
|
||||
self._retrievers.clear()
|
||||
|
||||
|
||||
def get_retriever_registry() -> RetrieverRegistry:
|
||||
"""Get the global retriever registry instance."""
|
||||
return RetrieverRegistry()
|
||||
|
||||
|
||||
def reciprocal_rank_fusion(
|
||||
result_lists: List[List[RetrievalResult]],
|
||||
k: int = 60
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
Reciprocal Rank Fusion (RRF) algorithm for combining multiple result lists.
|
||||
|
||||
Pure Python implementation with no external dependencies.
|
||||
|
||||
Args:
|
||||
result_lists: List of result lists from different retrievers
|
||||
k: RRF constant (default 60)
|
||||
|
||||
Returns:
|
||||
Fused and re-ranked list of results
|
||||
"""
|
||||
# Track scores by document ID
|
||||
fused_scores: Dict[str, float] = {}
|
||||
doc_map: Dict[str, RetrievalResult] = {}
|
||||
|
||||
for results in result_lists:
|
||||
for rank, result in enumerate(results):
|
||||
doc_key = result.doc_id or result.text[:100] # Use text prefix as fallback key
|
||||
|
||||
# RRF score: 1 / (k + rank)
|
||||
rrf_score = 1.0 / (k + rank + 1)
|
||||
|
||||
if doc_key in fused_scores:
|
||||
fused_scores[doc_key] += rrf_score
|
||||
else:
|
||||
fused_scores[doc_key] = rrf_score
|
||||
doc_map[doc_key] = result
|
||||
|
||||
# Sort by fused score
|
||||
sorted_keys = sorted(fused_scores.keys(), key=lambda x: fused_scores[x], reverse=True)
|
||||
|
||||
# Build result list with updated scores
|
||||
fused_results = []
|
||||
for key in sorted_keys:
|
||||
result = doc_map[key]
|
||||
fused_results.append(RetrievalResult(
|
||||
text=result.text,
|
||||
score=fused_scores[key],
|
||||
metadata=result.metadata,
|
||||
doc_id=result.doc_id,
|
||||
chunk_index=result.chunk_index
|
||||
))
|
||||
|
||||
return fused_results
|
||||
|
||||
|
||||
def merge_adjacent_chunks(
|
||||
results: List[RetrievalResult],
|
||||
max_gap: int = 1
|
||||
) -> List[RetrievalResult]:
|
||||
"""
|
||||
Merge adjacent chunks from the same document.
|
||||
|
||||
Args:
|
||||
results: List of retrieval results
|
||||
max_gap: Maximum gap between chunk indices to merge
|
||||
|
||||
Returns:
|
||||
Merged results
|
||||
"""
|
||||
if not results:
|
||||
return results
|
||||
|
||||
# Group by document ID
|
||||
doc_groups: Dict[str, List[RetrievalResult]] = {}
|
||||
for result in results:
|
||||
doc_id = result.doc_id or "unknown"
|
||||
if doc_id not in doc_groups:
|
||||
doc_groups[doc_id] = []
|
||||
doc_groups[doc_id].append(result)
|
||||
|
||||
merged_results = []
|
||||
|
||||
for doc_id, chunks in doc_groups.items():
|
||||
# Sort by chunk index
|
||||
chunks_with_index = [c for c in chunks if c.chunk_index is not None]
|
||||
chunks_without_index = [c for c in chunks if c.chunk_index is None]
|
||||
|
||||
if not chunks_with_index:
|
||||
merged_results.extend(chunks_without_index)
|
||||
continue
|
||||
|
||||
chunks_with_index.sort(key=lambda x: x.chunk_index or 0)
|
||||
|
||||
# Merge adjacent chunks
|
||||
current_group = [chunks_with_index[0]]
|
||||
|
||||
for chunk in chunks_with_index[1:]:
|
||||
last_chunk = current_group[-1]
|
||||
last_idx = last_chunk.chunk_index or 0
|
||||
curr_idx = chunk.chunk_index or 0
|
||||
|
||||
if curr_idx - last_idx <= max_gap:
|
||||
current_group.append(chunk)
|
||||
else:
|
||||
# Merge current group and start new one
|
||||
merged_results.append(_merge_chunk_group(current_group))
|
||||
current_group = [chunk]
|
||||
|
||||
# Don't forget the last group
|
||||
if current_group:
|
||||
merged_results.append(_merge_chunk_group(current_group))
|
||||
|
||||
merged_results.extend(chunks_without_index)
|
||||
|
||||
# Sort by score
|
||||
merged_results.sort(key=lambda x: x.score, reverse=True)
|
||||
return merged_results
|
||||
|
||||
|
||||
def _merge_chunk_group(chunks: List[RetrievalResult]) -> RetrievalResult:
|
||||
"""Merge a group of chunks into a single result."""
|
||||
if len(chunks) == 1:
|
||||
return chunks[0]
|
||||
|
||||
# Combine text
|
||||
merged_text = "\n".join(c.text for c in chunks)
|
||||
|
||||
# Average score
|
||||
avg_score = sum(c.score for c in chunks) / len(chunks)
|
||||
|
||||
# Use first chunk's metadata
|
||||
metadata = chunks[0].metadata.copy()
|
||||
metadata["merged_chunks"] = len(chunks)
|
||||
|
||||
return RetrievalResult(
|
||||
text=merged_text,
|
||||
score=avg_score,
|
||||
metadata=metadata,
|
||||
doc_id=chunks[0].doc_id,
|
||||
chunk_index=chunks[0].chunk_index
|
||||
)
|
||||
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
Vector Store Adapter Protocol for PraisonAI Agents.
|
||||
|
||||
This module defines the protocol for vector stores and a registry system.
|
||||
NO heavy imports - only stdlib and typing.
|
||||
|
||||
Implementations are provided by praisonai-tools or wrapper layer.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorRecord:
|
||||
"""
|
||||
A record in a vector store.
|
||||
|
||||
Attributes:
|
||||
id: Unique identifier
|
||||
text: The text content
|
||||
embedding: The vector embedding
|
||||
metadata: Optional metadata dict
|
||||
score: Optional similarity score (for query results)
|
||||
"""
|
||||
id: str
|
||||
text: str
|
||||
embedding: List[float]
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
score: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"text": self.text,
|
||||
"embedding": self.embedding,
|
||||
"metadata": self.metadata,
|
||||
"score": self.score
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "VectorRecord":
|
||||
"""Create from dictionary."""
|
||||
return cls(
|
||||
id=data["id"],
|
||||
text=data["text"],
|
||||
embedding=data["embedding"],
|
||||
metadata=data.get("metadata", {}),
|
||||
score=data.get("score")
|
||||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VectorStoreProtocol(Protocol):
|
||||
"""
|
||||
Protocol for vector stores.
|
||||
|
||||
Implementations must provide methods for:
|
||||
- Adding vectors
|
||||
- Querying by vector similarity
|
||||
- Deleting vectors
|
||||
- Getting count
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def add(
|
||||
self,
|
||||
texts: List[str],
|
||||
embeddings: List[List[float]],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
Add vectors to the store.
|
||||
|
||||
Args:
|
||||
texts: List of text content
|
||||
embeddings: List of embedding vectors
|
||||
metadatas: Optional list of metadata dicts
|
||||
ids: Optional list of IDs (auto-generated if not provided)
|
||||
namespace: Optional namespace for multi-tenant isolation
|
||||
|
||||
Returns:
|
||||
List of IDs for the added vectors
|
||||
"""
|
||||
...
|
||||
|
||||
def query(
|
||||
self,
|
||||
embedding: List[float],
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""
|
||||
Query vectors by similarity.
|
||||
|
||||
Args:
|
||||
embedding: Query embedding vector
|
||||
top_k: Number of results to return
|
||||
namespace: Optional namespace filter
|
||||
filter: Optional metadata filter
|
||||
|
||||
Returns:
|
||||
List of VectorRecord objects sorted by similarity
|
||||
"""
|
||||
...
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""
|
||||
Delete vectors from the store.
|
||||
|
||||
Args:
|
||||
ids: Optional list of IDs to delete
|
||||
namespace: Optional namespace filter
|
||||
filter: Optional metadata filter
|
||||
delete_all: If True, delete all vectors
|
||||
|
||||
Returns:
|
||||
Number of vectors deleted
|
||||
"""
|
||||
...
|
||||
|
||||
def count(self, namespace: Optional[str] = None) -> int:
|
||||
"""
|
||||
Get the number of vectors in the store.
|
||||
|
||||
Args:
|
||||
namespace: Optional namespace filter
|
||||
|
||||
Returns:
|
||||
Number of vectors
|
||||
"""
|
||||
...
|
||||
|
||||
def get(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""
|
||||
Get vectors by ID.
|
||||
|
||||
Args:
|
||||
ids: List of IDs to retrieve
|
||||
namespace: Optional namespace filter
|
||||
|
||||
Returns:
|
||||
List of VectorRecord objects
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class VectorStoreRegistry:
|
||||
"""
|
||||
Registry for vector stores.
|
||||
|
||||
Provides lazy loading and configuration-based initialization.
|
||||
"""
|
||||
|
||||
_instance: Optional["VectorStoreRegistry"] = None
|
||||
|
||||
def __new__(cls) -> "VectorStoreRegistry":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._stores: Dict[str, Callable[..., VectorStoreProtocol]] = {}
|
||||
cls._instance._initialized_stores: Dict[str, VectorStoreProtocol] = {}
|
||||
return cls._instance
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
factory: Callable[..., VectorStoreProtocol]
|
||||
) -> None:
|
||||
"""
|
||||
Register a vector store factory.
|
||||
|
||||
Args:
|
||||
name: Unique name for the store (e.g., "chroma", "pinecone")
|
||||
factory: Callable that returns a VectorStoreProtocol instance
|
||||
"""
|
||||
self._stores[name] = factory
|
||||
|
||||
def get(
|
||||
self,
|
||||
name: str,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> Optional[VectorStoreProtocol]:
|
||||
"""
|
||||
Get a vector store by name (lazy initialization).
|
||||
|
||||
Args:
|
||||
name: Store name
|
||||
config: Optional configuration dict
|
||||
namespace: Optional namespace for multi-tenant isolation
|
||||
|
||||
Returns:
|
||||
VectorStoreProtocol instance or None
|
||||
"""
|
||||
cache_key = f"{name}:{namespace or 'default'}"
|
||||
|
||||
if cache_key in self._initialized_stores:
|
||||
return self._initialized_stores[cache_key]
|
||||
|
||||
if name in self._stores:
|
||||
try:
|
||||
store = self._stores[name](config=config, namespace=namespace)
|
||||
self._initialized_stores[cache_key] = store
|
||||
return store
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize vector store '{name}': {e}")
|
||||
return None
|
||||
return None
|
||||
|
||||
def list_stores(self) -> List[str]:
|
||||
"""List all registered store names."""
|
||||
return list(self._stores.keys())
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all registered stores (mainly for testing)."""
|
||||
self._stores.clear()
|
||||
self._initialized_stores.clear()
|
||||
|
||||
|
||||
def get_vector_store_registry() -> VectorStoreRegistry:
|
||||
"""Get the global vector store registry instance."""
|
||||
return VectorStoreRegistry()
|
||||
|
||||
|
||||
class InMemoryVectorStore:
|
||||
"""
|
||||
Pure Python in-memory vector store for testing and default usage.
|
||||
|
||||
Uses cosine similarity for queries. No external dependencies.
|
||||
"""
|
||||
|
||||
name: str = "memory"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
namespace: Optional[str] = None
|
||||
):
|
||||
self.config = config or {}
|
||||
self.default_namespace = namespace or "default"
|
||||
self._data: Dict[str, Dict[str, VectorRecord]] = {} # namespace -> id -> record
|
||||
|
||||
def _get_namespace_data(self, namespace: Optional[str] = None) -> Dict[str, VectorRecord]:
|
||||
"""Get or create namespace data."""
|
||||
ns = namespace or self.default_namespace
|
||||
if ns not in self._data:
|
||||
self._data[ns] = {}
|
||||
return self._data[ns]
|
||||
|
||||
def add(
|
||||
self,
|
||||
texts: List[str],
|
||||
embeddings: List[List[float]],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""Add vectors to the store."""
|
||||
import uuid
|
||||
|
||||
ns_data = self._get_namespace_data(namespace)
|
||||
metadatas = metadatas or [{} for _ in texts]
|
||||
ids = ids or [str(uuid.uuid4()) for _ in texts]
|
||||
|
||||
result_ids = []
|
||||
for i, (text, embedding, metadata, id_) in enumerate(zip(texts, embeddings, metadatas, ids)):
|
||||
record = VectorRecord(
|
||||
id=id_,
|
||||
text=text,
|
||||
embedding=embedding,
|
||||
metadata=metadata
|
||||
)
|
||||
ns_data[id_] = record
|
||||
result_ids.append(id_)
|
||||
|
||||
return result_ids
|
||||
|
||||
def query(
|
||||
self,
|
||||
embedding: List[float],
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""Query vectors by cosine similarity."""
|
||||
ns_data = self._get_namespace_data(namespace)
|
||||
|
||||
if not ns_data:
|
||||
return []
|
||||
|
||||
# Calculate cosine similarity for all records
|
||||
results = []
|
||||
for record in ns_data.values():
|
||||
# Apply metadata filter if provided
|
||||
if filter:
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if record.metadata.get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if not match:
|
||||
continue
|
||||
|
||||
score = self._cosine_similarity(embedding, record.embedding)
|
||||
results.append(VectorRecord(
|
||||
id=record.id,
|
||||
text=record.text,
|
||||
embedding=record.embedding,
|
||||
metadata=record.metadata,
|
||||
score=score
|
||||
))
|
||||
|
||||
# Sort by score descending and return top_k
|
||||
results.sort(key=lambda x: x.score or 0, reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete vectors from the store."""
|
||||
ns = namespace or self.default_namespace
|
||||
|
||||
if delete_all:
|
||||
if ns in self._data:
|
||||
count = len(self._data[ns])
|
||||
self._data[ns] = {}
|
||||
return count
|
||||
return 0
|
||||
|
||||
if ns not in self._data:
|
||||
return 0
|
||||
|
||||
ns_data = self._data[ns]
|
||||
deleted = 0
|
||||
|
||||
if ids:
|
||||
for id_ in ids:
|
||||
if id_ in ns_data:
|
||||
del ns_data[id_]
|
||||
deleted += 1
|
||||
elif filter:
|
||||
to_delete = []
|
||||
for id_, record in ns_data.items():
|
||||
match = True
|
||||
for key, value in filter.items():
|
||||
if record.metadata.get(key) != value:
|
||||
match = False
|
||||
break
|
||||
if match:
|
||||
to_delete.append(id_)
|
||||
for id_ in to_delete:
|
||||
del ns_data[id_]
|
||||
deleted += 1
|
||||
|
||||
return deleted
|
||||
|
||||
def count(self, namespace: Optional[str] = None) -> int:
|
||||
"""Get the number of vectors in the store."""
|
||||
ns = namespace or self.default_namespace
|
||||
if ns in self._data:
|
||||
return len(self._data[ns])
|
||||
return 0
|
||||
|
||||
def get(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None
|
||||
) -> List[VectorRecord]:
|
||||
"""Get vectors by ID."""
|
||||
ns_data = self._get_namespace_data(namespace)
|
||||
results = []
|
||||
for id_ in ids:
|
||||
if id_ in ns_data:
|
||||
results.append(ns_data[id_])
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def _cosine_similarity(a: List[float], b: List[float]) -> float:
|
||||
"""Calculate cosine similarity between two vectors."""
|
||||
if len(a) != len(b):
|
||||
return 0.0
|
||||
|
||||
dot_product = sum(x * y for x, y in zip(a, b))
|
||||
norm_a = sum(x * x for x in a) ** 0.5
|
||||
norm_b = sum(x * x for x in b) ** 0.5
|
||||
|
||||
if norm_a == 0 or norm_b == 0:
|
||||
return 0.0
|
||||
|
||||
return dot_product / (norm_a * norm_b)
|
||||
|
||||
|
||||
# Register the in-memory store by default
|
||||
def _register_default_stores():
|
||||
"""Register default vector stores."""
|
||||
registry = get_vector_store_registry()
|
||||
registry.register("memory", InMemoryVectorStore)
|
||||
|
||||
|
||||
_register_default_stores()
|
||||
@@ -4,16 +4,20 @@ Session Management for PraisonAI Agents
|
||||
A simple wrapper around existing stateful capabilities to provide a unified
|
||||
session API for developers building stateful agent applications.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import uuid
|
||||
import requests
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory.memory import Memory
|
||||
from .knowledge.knowledge import Knowledge
|
||||
|
||||
from .agent import Agent
|
||||
from .memory import Memory
|
||||
from .knowledge import Knowledge
|
||||
|
||||
|
||||
class Session:
|
||||
@@ -112,6 +116,7 @@ class Session:
|
||||
if self.is_remote:
|
||||
raise ValueError("Memory operations are not available for remote agent sessions")
|
||||
if self._memory is None:
|
||||
from .memory.memory import Memory
|
||||
self._memory = Memory(config=self.memory_config)
|
||||
return self._memory
|
||||
|
||||
@@ -121,6 +126,7 @@ class Session:
|
||||
if self.is_remote:
|
||||
raise ValueError("Knowledge operations are not available for remote agent sessions")
|
||||
if self._knowledge is None:
|
||||
from .knowledge.knowledge import Knowledge
|
||||
self._knowledge = Knowledge(config=self.knowledge_config)
|
||||
return self._knowledge
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.3.0"
|
||||
version = "0.3.1"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
Generated
+1
-1
@@ -2422,7 +2422,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.3.0"
|
||||
version = "0.3.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -204,6 +204,12 @@ npm install praisonai
|
||||
| ↳ Add Custom Knowledge | [Example](examples/python/concepts/knowledge-agents.py) | [📖](https://docs.praison.ai/features/knowledge) |
|
||||
| ↳ RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [📖](https://docs.praison.ai/features/rag) |
|
||||
| ↳ Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [📖](https://docs.praison.ai/features/chat-with-pdf) |
|
||||
| ↳ Data Readers (PDF, DOCX, etc.) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-readers-api) |
|
||||
| ↳ Vector Store Selection | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-vector-store-api) |
|
||||
| ↳ Retrieval Strategies | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-retrieval-api) |
|
||||
| ↳ Rerankers | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-reranker-api) |
|
||||
| ↳ Index Types (Vector/Keyword/Hybrid) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-index-api) |
|
||||
| ↳ Query Engines (Sub-Question, etc.) | [CLI](#knowledge-cli) | [📖](https://docs.praison.ai/api/praisonai/knowledge-query-engine-api) |
|
||||
| **🔬 Research & Intelligence** | | |
|
||||
| ↳ Deep Research Agents | [Example](examples/python/agents/research-agent.py) | [📖](https://docs.praison.ai/agents/deep-research) |
|
||||
| ↳ Query Rewriter Agent | [Example](#5-query-rewriter-agent) | [📖](https://docs.praison.ai/agents/query-rewriter) |
|
||||
@@ -2641,6 +2647,84 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
|
||||
---
|
||||
|
||||
## 📚 Knowledge & Retrieval (RAG)
|
||||
|
||||
PraisonAI provides a complete knowledge stack for building RAG applications with multiple vector stores, retrieval strategies, rerankers, and query modes.
|
||||
|
||||
### Knowledge CLI Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai knowledge add <file\|dir\|url>` | Add documents to knowledge base |
|
||||
| `praisonai knowledge query <question>` | Query knowledge base with RAG |
|
||||
| `praisonai knowledge list` | List indexed documents |
|
||||
| `praisonai knowledge clear` | Clear knowledge base |
|
||||
| `praisonai knowledge stats` | Show knowledge base statistics |
|
||||
|
||||
### Knowledge CLI Options
|
||||
|
||||
| Option | Values | Description |
|
||||
|--------|--------|-------------|
|
||||
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | Vector store backend |
|
||||
| `--retrieval` | `basic`, `fusion`, `recursive`, `auto_merge` | Retrieval strategy |
|
||||
| `--reranker` | `simple`, `llm`, `cross_encoder`, `cohere` | Reranking method |
|
||||
| `--index-type` | `vector`, `keyword`, `hybrid` | Index type |
|
||||
| `--query-mode` | `default`, `sub_question`, `summarize` | Query mode |
|
||||
|
||||
### Knowledge CLI Examples
|
||||
|
||||
```bash
|
||||
# Add documents
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge add https://example.com/page.html
|
||||
praisonai knowledge add "*.pdf"
|
||||
|
||||
# Query with advanced options
|
||||
praisonai knowledge query "How to authenticate?" --retrieval fusion --reranker llm
|
||||
|
||||
# Full advanced query
|
||||
praisonai knowledge query "authentication flow" \
|
||||
--vector-store chroma \
|
||||
--retrieval fusion \
|
||||
--reranker llm \
|
||||
--index-type hybrid \
|
||||
--query-mode sub_question
|
||||
```
|
||||
|
||||
### Knowledge SDK Usage
|
||||
|
||||
```python
|
||||
from praisonaiagents import Agent, Knowledge
|
||||
|
||||
# Simple usage with Agent
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
knowledge=["docs/manual.pdf", "data/faq.txt"],
|
||||
knowledge_config={
|
||||
"vector_store": {"provider": "chroma"}
|
||||
}
|
||||
)
|
||||
response = agent.chat("How do I authenticate?")
|
||||
|
||||
# Direct Knowledge usage
|
||||
knowledge = Knowledge()
|
||||
knowledge.add("document.pdf")
|
||||
results = knowledge.search("authentication", limit=5)
|
||||
```
|
||||
|
||||
### Knowledge Stack Features Table
|
||||
|
||||
| Feature | Description | SDK Docs | CLI Docs |
|
||||
|---------|-------------|----------|----------|
|
||||
| **Data Readers** | Load PDF, Markdown, Text, HTML, URLs | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Vector Stores** | ChromaDB, Pinecone, Qdrant, Weaviate, In-Memory | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Retrieval Strategies** | Basic, Fusion (RRF), Recursive, Auto-Merge | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Rerankers** | Simple, LLM, Cross-Encoder, Cohere | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Index Types** | Vector, Keyword (BM25), Hybrid | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
| **Query Engines** | Default, Sub-Question, Summarize | [SDK](/docs/sdk/praisonaiagents/knowledge/protocols) | [CLI](/docs/cli/knowledge) |
|
||||
|
||||
---
|
||||
|
||||
## 🔬 Advanced Features
|
||||
|
||||
### Research & Intelligence
|
||||
@@ -2755,6 +2839,44 @@ agent.chat("Hello!") # Auto-persists messages, runs, traces
|
||||
| `praisonai persistence migrate` | Apply schema migrations |
|
||||
| `praisonai persistence status` | Show schema status |
|
||||
|
||||
### Knowledge CLI Commands {#knowledge-cli}
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `praisonai knowledge add <source>` | Add file, directory, URL, or glob pattern |
|
||||
| `praisonai knowledge query "<question>"` | Query knowledge base with RAG |
|
||||
| `praisonai knowledge list` | List indexed documents |
|
||||
| `praisonai knowledge clear` | Clear knowledge base |
|
||||
| `praisonai knowledge stats` | Show knowledge base statistics |
|
||||
|
||||
**Knowledge Query Flags:**
|
||||
|
||||
| Flag | Values | Default |
|
||||
|------|--------|---------|
|
||||
| `--vector-store` | `memory`, `chroma`, `pinecone`, `qdrant`, `weaviate` | `chroma` |
|
||||
| `--retrieval-strategy` | `basic`, `fusion`, `recursive`, `auto_merge` | `basic` |
|
||||
| `--reranker` | `none`, `simple`, `llm`, `cross_encoder`, `cohere` | `none` |
|
||||
| `--index-type` | `vector`, `keyword`, `hybrid` | `vector` |
|
||||
| `--query-mode` | `default`, `sub_question`, `summarize` | `default` |
|
||||
| `--workspace` | Path to workspace directory | Current dir |
|
||||
| `--session` | Session ID for persistence | - |
|
||||
|
||||
**Examples:**
|
||||
|
||||
```bash
|
||||
# Add documents
|
||||
praisonai knowledge add document.pdf
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge add "*.md"
|
||||
|
||||
# Query with options
|
||||
praisonai knowledge query "How to authenticate?" \
|
||||
--vector-store chroma \
|
||||
--retrieval-strategy fusion \
|
||||
--reranker simple \
|
||||
--query-mode sub_question
|
||||
```
|
||||
|
||||
### Databases Table
|
||||
|
||||
| Database | Store Type | Install | Example | Docs |
|
||||
|
||||
@@ -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.5.0.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.5.0.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.5.1.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.5.1.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
PraisonAI Adapters - Implementations for core protocols.
|
||||
|
||||
This module provides concrete implementations of:
|
||||
- Reader adapters (AutoReader, LlamaIndexReaderAdapter, MarkItDownReaderAdapter)
|
||||
- Vector store adapters (ChromaAdapter, PineconeAdapter, etc.)
|
||||
- Retriever implementations
|
||||
- Reranker implementations
|
||||
"""
|
||||
|
||||
# Lazy loading to avoid heavy imports at package load time
|
||||
_LAZY_IMPORTS = {
|
||||
# Readers
|
||||
"AutoReader": ("praisonai.adapters.readers", "AutoReader"),
|
||||
"MarkItDownReader": ("praisonai.adapters.readers", "MarkItDownReader"),
|
||||
"TextReader": ("praisonai.adapters.readers", "TextReader"),
|
||||
"DirectoryReader": ("praisonai.adapters.readers", "DirectoryReader"),
|
||||
"register_default_readers": ("praisonai.adapters.readers", "register_default_readers"),
|
||||
|
||||
# Vector stores
|
||||
"ChromaVectorStore": ("praisonai.adapters.vector_stores", "ChromaVectorStore"),
|
||||
"register_default_vector_stores": ("praisonai.adapters.vector_stores", "register_default_vector_stores"),
|
||||
|
||||
# Retrievers
|
||||
"BasicRetriever": ("praisonai.adapters.retrievers", "BasicRetriever"),
|
||||
"FusionRetriever": ("praisonai.adapters.retrievers", "FusionRetriever"),
|
||||
"register_default_retrievers": ("praisonai.adapters.retrievers", "register_default_retrievers"),
|
||||
|
||||
# Rerankers
|
||||
"LLMReranker": ("praisonai.adapters.rerankers", "LLMReranker"),
|
||||
"register_default_rerankers": ("praisonai.adapters.rerankers", "register_default_rerankers"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load adapters."""
|
||||
if name in _LAZY_IMPORTS:
|
||||
module_path, attr_name = _LAZY_IMPORTS[name]
|
||||
import importlib
|
||||
module = importlib.import_module(module_path)
|
||||
return getattr(module, attr_name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__():
|
||||
"""List available attributes."""
|
||||
return list(_LAZY_IMPORTS.keys())
|
||||
|
||||
|
||||
__all__ = list(_LAZY_IMPORTS.keys())
|
||||
@@ -0,0 +1,395 @@
|
||||
"""
|
||||
Reader Adapters for PraisonAI.
|
||||
|
||||
Provides concrete implementations of ReaderProtocol:
|
||||
- TextReader: Plain text files
|
||||
- MarkItDownReader: Uses markitdown for documents
|
||||
- DirectoryReader: Reads all files in a directory
|
||||
- AutoReader: Automatic file type detection and routing
|
||||
"""
|
||||
|
||||
import os
|
||||
import glob as glob_module
|
||||
import logging
|
||||
import importlib.util
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy import flags
|
||||
_MARKITDOWN_AVAILABLE = None
|
||||
_LLAMA_INDEX_AVAILABLE = None
|
||||
|
||||
|
||||
def _check_markitdown():
|
||||
"""Check if markitdown is available."""
|
||||
global _MARKITDOWN_AVAILABLE
|
||||
if _MARKITDOWN_AVAILABLE is None:
|
||||
_MARKITDOWN_AVAILABLE = importlib.util.find_spec("markitdown") is not None
|
||||
return _MARKITDOWN_AVAILABLE
|
||||
|
||||
|
||||
def _check_llama_index():
|
||||
"""Check if llama_index is available."""
|
||||
global _LLAMA_INDEX_AVAILABLE
|
||||
if _LLAMA_INDEX_AVAILABLE is None:
|
||||
_LLAMA_INDEX_AVAILABLE = importlib.util.find_spec("llama_index") is not None
|
||||
return _LLAMA_INDEX_AVAILABLE
|
||||
|
||||
|
||||
class TextReader:
|
||||
"""Simple text file reader."""
|
||||
|
||||
name: str = "text"
|
||||
supported_extensions: List[str] = ["txt", "text", "log"]
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load a text file."""
|
||||
from praisonaiagents.knowledge.readers import Document
|
||||
|
||||
try:
|
||||
with open(source, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
|
||||
doc_metadata = metadata or {}
|
||||
doc_metadata.update({
|
||||
"source": source,
|
||||
"filename": os.path.basename(source),
|
||||
"file_type": "text"
|
||||
})
|
||||
|
||||
return [Document(content=content, metadata=doc_metadata)]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to read text file {source}: {e}")
|
||||
return []
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""Check if this reader can handle the source."""
|
||||
ext = os.path.splitext(source)[1].lower().lstrip(".")
|
||||
return ext in self.supported_extensions
|
||||
|
||||
|
||||
class MarkItDownReader:
|
||||
"""Reader using markitdown for document conversion."""
|
||||
|
||||
name: str = "markitdown"
|
||||
supported_extensions: List[str] = [
|
||||
"pdf", "doc", "docx", "ppt", "pptx", "xls", "xlsx",
|
||||
"html", "htm", "md", "markdown", "csv", "json", "xml",
|
||||
"jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp",
|
||||
"mp3", "wav", "ogg", "m4a", "flac"
|
||||
]
|
||||
|
||||
def __init__(self):
|
||||
self._converter = None
|
||||
|
||||
@property
|
||||
def converter(self):
|
||||
"""Lazy load markitdown converter."""
|
||||
if self._converter is None:
|
||||
if not _check_markitdown():
|
||||
raise ImportError(
|
||||
"markitdown is required for document conversion. "
|
||||
"Install with: pip install markitdown"
|
||||
)
|
||||
from markitdown import MarkItDown
|
||||
self._converter = MarkItDown()
|
||||
return self._converter
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load a document using markitdown."""
|
||||
from praisonaiagents.knowledge.readers import Document
|
||||
|
||||
try:
|
||||
result = self.converter.convert(source)
|
||||
content = result.text_content
|
||||
|
||||
if not content:
|
||||
logger.warning(f"No content extracted from {source}")
|
||||
return []
|
||||
|
||||
doc_metadata = metadata or {}
|
||||
doc_metadata.update({
|
||||
"source": source,
|
||||
"filename": os.path.basename(source),
|
||||
"file_type": os.path.splitext(source)[1].lower().lstrip(".")
|
||||
})
|
||||
|
||||
return [Document(content=content, metadata=doc_metadata)]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to convert document {source}: {e}")
|
||||
return []
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""Check if this reader can handle the source."""
|
||||
ext = os.path.splitext(source)[1].lower().lstrip(".")
|
||||
return ext in self.supported_extensions and _check_markitdown()
|
||||
|
||||
|
||||
class DirectoryReader:
|
||||
"""Reader for directories - recursively reads all files."""
|
||||
|
||||
name: str = "directory"
|
||||
supported_extensions: List[str] = [] # Handles directories, not extensions
|
||||
|
||||
def __init__(self, recursive: bool = True, exclude_patterns: Optional[List[str]] = None):
|
||||
self.recursive = recursive
|
||||
self.exclude_patterns = exclude_patterns or [
|
||||
"*.pyc", "__pycache__", ".git", ".svn", "node_modules",
|
||||
"*.egg-info", ".env", ".venv", "venv"
|
||||
]
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load all files from a directory."""
|
||||
from praisonaiagents.knowledge.readers import get_reader_registry
|
||||
|
||||
if not os.path.isdir(source):
|
||||
logger.error(f"Not a directory: {source}")
|
||||
return []
|
||||
|
||||
documents = []
|
||||
registry = get_reader_registry()
|
||||
|
||||
# Walk directory
|
||||
if self.recursive:
|
||||
for root, dirs, files in os.walk(source):
|
||||
# Filter out excluded directories
|
||||
dirs[:] = [d for d in dirs if not self._should_exclude(d)]
|
||||
|
||||
for file in files:
|
||||
if self._should_exclude(file):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(root, file)
|
||||
reader = registry.get_for_source(file_path)
|
||||
|
||||
if reader and reader.name != "directory":
|
||||
file_metadata = (metadata or {}).copy()
|
||||
file_metadata["parent_dir"] = source
|
||||
loaded_docs = reader.load(file_path, metadata=file_metadata)
|
||||
documents.extend(loaded_docs)
|
||||
else:
|
||||
for file in os.listdir(source):
|
||||
if self._should_exclude(file):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(source, file)
|
||||
if os.path.isfile(file_path):
|
||||
reader = registry.get_for_source(file_path)
|
||||
if reader and reader.name != "directory":
|
||||
file_metadata = (metadata or {}).copy()
|
||||
file_metadata["parent_dir"] = source
|
||||
docs = reader.load(file_path, metadata=file_metadata)
|
||||
documents.extend(docs)
|
||||
|
||||
return documents
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""Check if this reader can handle the source."""
|
||||
return os.path.isdir(source)
|
||||
|
||||
def _should_exclude(self, name: str) -> bool:
|
||||
"""Check if a file/directory should be excluded."""
|
||||
import fnmatch
|
||||
for pattern in self.exclude_patterns:
|
||||
if fnmatch.fnmatch(name, pattern):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class GlobReader:
|
||||
"""Reader for glob patterns."""
|
||||
|
||||
name: str = "glob"
|
||||
supported_extensions: List[str] = []
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load files matching a glob pattern."""
|
||||
from praisonaiagents.knowledge.readers import get_reader_registry
|
||||
|
||||
documents = []
|
||||
registry = get_reader_registry()
|
||||
|
||||
for file_path in glob_module.glob(source, recursive=True):
|
||||
if os.path.isfile(file_path):
|
||||
reader = registry.get_for_source(file_path)
|
||||
if reader and reader.name not in ("glob", "directory"):
|
||||
file_metadata = (metadata or {}).copy()
|
||||
file_metadata["glob_pattern"] = source
|
||||
docs = reader.load(file_path, metadata=file_metadata)
|
||||
documents.extend(docs)
|
||||
|
||||
return documents
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""Check if this is a glob pattern."""
|
||||
return "*" in source or "?" in source or "[" in source
|
||||
|
||||
|
||||
class URLReader:
|
||||
"""Reader for URLs (HTML pages)."""
|
||||
|
||||
name: str = "url"
|
||||
supported_extensions: List[str] = []
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load content from a URL."""
|
||||
from praisonaiagents.knowledge.readers import Document
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError:
|
||||
try:
|
||||
import requests as httpx
|
||||
except ImportError:
|
||||
logger.error("httpx or requests required for URL reading")
|
||||
return []
|
||||
|
||||
try:
|
||||
response = httpx.get(source, timeout=30, follow_redirects=True)
|
||||
response.raise_for_status()
|
||||
content = response.text
|
||||
|
||||
# Try to extract text from HTML
|
||||
try:
|
||||
from bs4 import BeautifulSoup
|
||||
soup = BeautifulSoup(content, 'html.parser')
|
||||
# Remove script and style elements
|
||||
for script in soup(["script", "style"]):
|
||||
script.decompose()
|
||||
content = soup.get_text(separator='\n', strip=True)
|
||||
except ImportError:
|
||||
# Fall back to raw content
|
||||
pass
|
||||
|
||||
doc_metadata = metadata or {}
|
||||
doc_metadata.update({
|
||||
"source": source,
|
||||
"url": source,
|
||||
"file_type": "html"
|
||||
})
|
||||
|
||||
return [Document(content=content, metadata=doc_metadata)]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch URL {source}: {e}")
|
||||
return []
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""Check if this is a URL."""
|
||||
return source.startswith(("http://", "https://"))
|
||||
|
||||
|
||||
class AutoReader:
|
||||
"""
|
||||
Automatic reader that detects source type and routes to appropriate reader.
|
||||
|
||||
Selection policy:
|
||||
1. If LlamaIndex reader is installed + supports the type → use it
|
||||
2. Else if MarkItDown supports → use it
|
||||
3. Else use simple built-in reader
|
||||
"""
|
||||
|
||||
name: str = "auto"
|
||||
supported_extensions: List[str] = [] # Handles everything
|
||||
|
||||
def load(
|
||||
self,
|
||||
source: str,
|
||||
*,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load from any source type."""
|
||||
from praisonaiagents.knowledge.readers import get_reader_registry, detect_source_kind
|
||||
|
||||
registry = get_reader_registry()
|
||||
source_kind = detect_source_kind(source)
|
||||
|
||||
# Route based on source kind
|
||||
if source_kind == "url":
|
||||
reader = registry.get("url")
|
||||
elif source_kind == "directory":
|
||||
reader = registry.get("directory")
|
||||
elif source_kind == "glob":
|
||||
reader = registry.get("glob")
|
||||
elif source_kind == "file":
|
||||
# Find best reader for file type
|
||||
reader = registry.get_for_source(source)
|
||||
if reader and reader.name == "auto":
|
||||
# Avoid infinite recursion - use markitdown or text
|
||||
ext = os.path.splitext(source)[1].lower().lstrip(".")
|
||||
if ext in MarkItDownReader.supported_extensions and _check_markitdown():
|
||||
reader = registry.get("markitdown")
|
||||
else:
|
||||
reader = registry.get("text")
|
||||
else:
|
||||
# Unknown - try markitdown then text
|
||||
if _check_markitdown():
|
||||
reader = registry.get("markitdown")
|
||||
else:
|
||||
reader = registry.get("text")
|
||||
|
||||
if reader:
|
||||
return reader.load(source, metadata=metadata)
|
||||
|
||||
logger.warning(f"No reader found for source: {source}")
|
||||
return []
|
||||
|
||||
def can_handle(self, source: str) -> bool:
|
||||
"""AutoReader can handle anything."""
|
||||
return True
|
||||
|
||||
|
||||
def register_default_readers():
|
||||
"""Register all default readers with the registry."""
|
||||
from praisonaiagents.knowledge.readers import get_reader_registry
|
||||
|
||||
registry = get_reader_registry()
|
||||
|
||||
# Register text reader
|
||||
registry.register("text", TextReader, ["txt", "text", "log"])
|
||||
|
||||
# Register markitdown reader
|
||||
registry.register("markitdown", MarkItDownReader, MarkItDownReader.supported_extensions)
|
||||
|
||||
# Register directory reader
|
||||
registry.register("directory", DirectoryReader)
|
||||
|
||||
# Register glob reader
|
||||
registry.register("glob", GlobReader)
|
||||
|
||||
# Register URL reader
|
||||
registry.register("url", URLReader)
|
||||
|
||||
# Register auto reader
|
||||
registry.register("auto", AutoReader)
|
||||
|
||||
logger.debug("Registered default readers")
|
||||
|
||||
|
||||
# Auto-register on import
|
||||
register_default_readers()
|
||||
@@ -0,0 +1,315 @@
|
||||
"""
|
||||
Reranker Implementations for PraisonAI.
|
||||
|
||||
Provides concrete implementations of RerankerProtocol:
|
||||
- LLMReranker: Uses LLM to rerank documents
|
||||
- CrossEncoderReranker: Uses sentence-transformers cross-encoder
|
||||
- CohereReranker: Uses Cohere rerank API
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy import flags
|
||||
_SENTENCE_TRANSFORMERS_AVAILABLE = None
|
||||
_COHERE_AVAILABLE = None
|
||||
|
||||
|
||||
def _check_sentence_transformers():
|
||||
"""Check if sentence-transformers is available."""
|
||||
global _SENTENCE_TRANSFORMERS_AVAILABLE
|
||||
if _SENTENCE_TRANSFORMERS_AVAILABLE is None:
|
||||
import importlib.util
|
||||
_SENTENCE_TRANSFORMERS_AVAILABLE = importlib.util.find_spec("sentence_transformers") is not None
|
||||
return _SENTENCE_TRANSFORMERS_AVAILABLE
|
||||
|
||||
|
||||
def _check_cohere():
|
||||
"""Check if cohere is available."""
|
||||
global _COHERE_AVAILABLE
|
||||
if _COHERE_AVAILABLE is None:
|
||||
import importlib.util
|
||||
_COHERE_AVAILABLE = importlib.util.find_spec("cohere") is not None
|
||||
return _COHERE_AVAILABLE
|
||||
|
||||
|
||||
class LLMReranker:
|
||||
"""
|
||||
LLM-based reranker.
|
||||
|
||||
Uses an LLM to score document relevance to a query.
|
||||
"""
|
||||
|
||||
name: str = "llm"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm: Optional[Any] = None,
|
||||
model: Optional[str] = None,
|
||||
batch_size: int = 5,
|
||||
**kwargs
|
||||
):
|
||||
self.llm = llm
|
||||
self.model = model or "gpt-4o-mini"
|
||||
self.batch_size = batch_size
|
||||
self._client = None
|
||||
|
||||
def _get_client(self):
|
||||
"""Lazy load OpenAI client."""
|
||||
if self._client is None:
|
||||
if self.llm:
|
||||
return self.llm
|
||||
import openai
|
||||
self._client = openai.OpenAI()
|
||||
return self._client
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Rerank documents using LLM scoring."""
|
||||
from praisonaiagents.knowledge.rerankers import RerankResult
|
||||
|
||||
results = []
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(documents), self.batch_size):
|
||||
batch = documents[i:i + self.batch_size]
|
||||
batch_results = self._score_batch(query, batch, start_idx=i)
|
||||
results.extend(batch_results)
|
||||
|
||||
# Sort by score
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
|
||||
if top_k is not None:
|
||||
results = results[:top_k]
|
||||
|
||||
return results
|
||||
|
||||
async def arerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.rerank(query, documents, top_k, **kwargs)
|
||||
|
||||
def _score_batch(self, query: str, documents: List[str], start_idx: int = 0) -> List[Any]:
|
||||
"""Score a batch of documents."""
|
||||
from praisonaiagents.knowledge.rerankers import RerankResult
|
||||
|
||||
results = []
|
||||
|
||||
for i, doc in enumerate(documents):
|
||||
score = self._score_document(query, doc)
|
||||
results.append(RerankResult(
|
||||
text=doc,
|
||||
score=score,
|
||||
original_index=start_idx + i
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
def _score_document(self, query: str, document: str) -> float:
|
||||
"""Score a single document's relevance to query."""
|
||||
try:
|
||||
client = self._get_client()
|
||||
|
||||
prompt = f"""Rate the relevance of this document to the query on a scale of 0 to 10.
|
||||
Return ONLY a number between 0 and 10.
|
||||
|
||||
Query: {query}
|
||||
|
||||
Document: {document[:1000]}
|
||||
|
||||
Relevance score (0-10):"""
|
||||
|
||||
if hasattr(client, 'chat'):
|
||||
# PraisonAI Agent or similar
|
||||
response = client.chat(prompt)
|
||||
else:
|
||||
# OpenAI client
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
max_tokens=10,
|
||||
temperature=0
|
||||
)
|
||||
response = response.choices[0].message.content
|
||||
|
||||
# Parse score
|
||||
if response:
|
||||
try:
|
||||
score = float(response.strip().split()[0])
|
||||
return min(max(score / 10.0, 0.0), 1.0) # Normalize to 0-1
|
||||
except (ValueError, IndexError):
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM scoring failed: {e}")
|
||||
|
||||
return 0.5 # Default score
|
||||
|
||||
|
||||
class CrossEncoderReranker:
|
||||
"""
|
||||
Cross-encoder reranker using sentence-transformers.
|
||||
|
||||
Uses a cross-encoder model for accurate relevance scoring.
|
||||
"""
|
||||
|
||||
name: str = "cross_encoder"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
**kwargs
|
||||
):
|
||||
if not _check_sentence_transformers():
|
||||
raise ImportError(
|
||||
"sentence-transformers is required for CrossEncoderReranker. "
|
||||
"Install with: pip install sentence-transformers"
|
||||
)
|
||||
|
||||
from sentence_transformers import CrossEncoder
|
||||
self.model = CrossEncoder(model_name)
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Rerank documents using cross-encoder."""
|
||||
from praisonaiagents.knowledge.rerankers import RerankResult
|
||||
|
||||
# Create query-document pairs
|
||||
pairs = [[query, doc] for doc in documents]
|
||||
|
||||
# Get scores
|
||||
scores = self.model.predict(pairs)
|
||||
|
||||
# Create results
|
||||
results = []
|
||||
for i, (doc, score) in enumerate(zip(documents, scores)):
|
||||
results.append(RerankResult(
|
||||
text=doc,
|
||||
score=float(score),
|
||||
original_index=i
|
||||
))
|
||||
|
||||
# Sort by score
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
|
||||
if top_k is not None:
|
||||
results = results[:top_k]
|
||||
|
||||
return results
|
||||
|
||||
async def arerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.rerank(query, documents, top_k, **kwargs)
|
||||
|
||||
|
||||
class CohereReranker:
|
||||
"""
|
||||
Cohere reranker using Cohere's rerank API.
|
||||
"""
|
||||
|
||||
name: str = "cohere"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
model: str = "rerank-english-v3.0",
|
||||
**kwargs
|
||||
):
|
||||
if not _check_cohere():
|
||||
raise ImportError(
|
||||
"cohere is required for CohereReranker. "
|
||||
"Install with: pip install cohere"
|
||||
)
|
||||
|
||||
import os
|
||||
import cohere
|
||||
|
||||
self.api_key = api_key or os.environ.get("COHERE_API_KEY")
|
||||
if not self.api_key:
|
||||
raise ValueError("COHERE_API_KEY environment variable or api_key parameter required")
|
||||
|
||||
self.model = model
|
||||
self._client = cohere.Client(self.api_key)
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Rerank documents using Cohere API."""
|
||||
from praisonaiagents.knowledge.rerankers import RerankResult
|
||||
|
||||
response = self._client.rerank(
|
||||
query=query,
|
||||
documents=documents,
|
||||
top_n=top_k or len(documents),
|
||||
model=self.model
|
||||
)
|
||||
|
||||
results = []
|
||||
for r in response.results:
|
||||
results.append(RerankResult(
|
||||
text=documents[r.index],
|
||||
score=r.relevance_score,
|
||||
original_index=r.index
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
async def arerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: List[str],
|
||||
top_k: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.rerank(query, documents, top_k, **kwargs)
|
||||
|
||||
|
||||
def register_default_rerankers():
|
||||
"""Register all default rerankers with the registry."""
|
||||
from praisonaiagents.knowledge.rerankers import get_reranker_registry
|
||||
|
||||
registry = get_reranker_registry()
|
||||
|
||||
# LLM reranker is always available
|
||||
registry.register("llm", LLMReranker)
|
||||
|
||||
# Cross-encoder if available
|
||||
if _check_sentence_transformers():
|
||||
registry.register("cross_encoder", CrossEncoderReranker)
|
||||
|
||||
# Cohere if available
|
||||
if _check_cohere():
|
||||
registry.register("cohere", CohereReranker)
|
||||
|
||||
logger.debug("Registered default rerankers")
|
||||
|
||||
|
||||
# Auto-register on import
|
||||
register_default_rerankers()
|
||||
@@ -0,0 +1,394 @@
|
||||
"""
|
||||
Retriever Implementations for PraisonAI.
|
||||
|
||||
Provides concrete implementations of RetrieverProtocol:
|
||||
- BasicRetriever: Simple vector similarity retrieval
|
||||
- FusionRetriever: Multi-query with RRF fusion
|
||||
- RecursiveRetriever: Depth-limited sub-retrieval
|
||||
- AutoMergeRetriever: Merges adjacent chunks
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BasicRetriever:
|
||||
"""
|
||||
Basic vector similarity retriever.
|
||||
|
||||
Simply queries the vector store and returns results.
|
||||
"""
|
||||
|
||||
name: str = "basic"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector_store: Any,
|
||||
embedding_fn: Optional[Callable[[str], List[float]]] = None,
|
||||
top_k: int = 10,
|
||||
**kwargs
|
||||
):
|
||||
self.vector_store = vector_store
|
||||
self.embedding_fn = embedding_fn
|
||||
self.top_k = top_k
|
||||
|
||||
def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Retrieve documents by vector similarity."""
|
||||
from praisonaiagents.knowledge.retrieval import RetrievalResult
|
||||
|
||||
k = top_k or self.top_k
|
||||
|
||||
# Get query embedding
|
||||
if self.embedding_fn:
|
||||
query_embedding = self.embedding_fn(query)
|
||||
else:
|
||||
raise ValueError("embedding_fn is required for BasicRetriever")
|
||||
|
||||
# Query vector store
|
||||
results = self.vector_store.query(
|
||||
embedding=query_embedding,
|
||||
top_k=k,
|
||||
filter=filter
|
||||
)
|
||||
|
||||
# Convert to RetrievalResult
|
||||
retrieval_results = []
|
||||
for r in results:
|
||||
retrieval_results.append(RetrievalResult(
|
||||
text=r.text if hasattr(r, 'text') else str(r),
|
||||
score=r.score if hasattr(r, 'score') else 1.0,
|
||||
metadata=r.metadata if hasattr(r, 'metadata') else {},
|
||||
doc_id=r.id if hasattr(r, 'id') else None,
|
||||
chunk_index=r.metadata.get('chunk_index') if hasattr(r, 'metadata') else None
|
||||
))
|
||||
|
||||
return retrieval_results
|
||||
|
||||
async def aretrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.retrieve(query, top_k, filter, **kwargs)
|
||||
|
||||
|
||||
class FusionRetriever:
|
||||
"""
|
||||
Fusion retriever using Reciprocal Rank Fusion.
|
||||
|
||||
Generates multiple query variations and fuses results.
|
||||
"""
|
||||
|
||||
name: str = "fusion"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector_store: Any,
|
||||
embedding_fn: Optional[Callable[[str], List[float]]] = None,
|
||||
llm: Optional[Any] = None,
|
||||
num_queries: int = 3,
|
||||
top_k: int = 10,
|
||||
rrf_k: int = 60,
|
||||
**kwargs
|
||||
):
|
||||
self.vector_store = vector_store
|
||||
self.embedding_fn = embedding_fn
|
||||
self.llm = llm
|
||||
self.num_queries = num_queries
|
||||
self.top_k = top_k
|
||||
self.rrf_k = rrf_k
|
||||
|
||||
def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Retrieve with query fusion."""
|
||||
from praisonaiagents.knowledge.retrieval import RetrievalResult, reciprocal_rank_fusion
|
||||
|
||||
k = top_k or self.top_k
|
||||
|
||||
# Generate query variations
|
||||
queries = self._generate_query_variations(query)
|
||||
|
||||
# Retrieve for each query
|
||||
all_results = []
|
||||
for q in queries:
|
||||
if self.embedding_fn:
|
||||
q_embedding = self.embedding_fn(q)
|
||||
results = self.vector_store.query(
|
||||
embedding=q_embedding,
|
||||
top_k=k,
|
||||
filter=filter
|
||||
)
|
||||
|
||||
# Convert to RetrievalResult
|
||||
query_results = []
|
||||
for r in results:
|
||||
query_results.append(RetrievalResult(
|
||||
text=r.text if hasattr(r, 'text') else str(r),
|
||||
score=r.score if hasattr(r, 'score') else 1.0,
|
||||
metadata=r.metadata if hasattr(r, 'metadata') else {},
|
||||
doc_id=r.id if hasattr(r, 'id') else None
|
||||
))
|
||||
all_results.append(query_results)
|
||||
|
||||
# Fuse results using RRF
|
||||
fused = reciprocal_rank_fusion(all_results, k=self.rrf_k)
|
||||
|
||||
return fused[:k]
|
||||
|
||||
async def aretrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.retrieve(query, top_k, filter, **kwargs)
|
||||
|
||||
def _generate_query_variations(self, query: str) -> List[str]:
|
||||
"""Generate query variations."""
|
||||
queries = [query]
|
||||
|
||||
# Try LLM-based variation if available
|
||||
if self.llm and len(queries) < self.num_queries:
|
||||
try:
|
||||
prompt = f"""Generate {self.num_queries - 1} alternative phrasings of this search query.
|
||||
Return only the queries, one per line.
|
||||
|
||||
Original query: {query}"""
|
||||
response = self.llm.chat(prompt)
|
||||
if response:
|
||||
for line in response.strip().split("\n"):
|
||||
line = line.strip()
|
||||
if line and line != query:
|
||||
queries.append(line)
|
||||
if len(queries) >= self.num_queries:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.warning(f"LLM query variation failed: {e}")
|
||||
|
||||
# Fallback: simple variations
|
||||
if len(queries) < self.num_queries:
|
||||
# Add question form
|
||||
if not query.endswith("?"):
|
||||
queries.append(f"What is {query}?")
|
||||
# Add keyword form
|
||||
keywords = " ".join(query.split()[:5])
|
||||
if keywords != query:
|
||||
queries.append(keywords)
|
||||
|
||||
return queries[:self.num_queries]
|
||||
|
||||
|
||||
class RecursiveRetriever:
|
||||
"""
|
||||
Recursive retriever with depth-limited sub-retrieval.
|
||||
|
||||
Retrieves, expands queries based on results, and retrieves again.
|
||||
"""
|
||||
|
||||
name: str = "recursive"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector_store: Any,
|
||||
embedding_fn: Optional[Callable[[str], List[float]]] = None,
|
||||
llm: Optional[Any] = None,
|
||||
max_depth: int = 2,
|
||||
top_k: int = 10,
|
||||
**kwargs
|
||||
):
|
||||
self.vector_store = vector_store
|
||||
self.embedding_fn = embedding_fn
|
||||
self.llm = llm
|
||||
self.max_depth = max_depth
|
||||
self.top_k = top_k
|
||||
|
||||
def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Retrieve with recursive expansion."""
|
||||
from praisonaiagents.knowledge.retrieval import RetrievalResult
|
||||
|
||||
k = top_k or self.top_k
|
||||
seen_ids = set()
|
||||
all_results = []
|
||||
|
||||
current_queries = [query]
|
||||
|
||||
for depth in range(self.max_depth):
|
||||
next_queries = []
|
||||
|
||||
for q in current_queries:
|
||||
if self.embedding_fn:
|
||||
q_embedding = self.embedding_fn(q)
|
||||
results = self.vector_store.query(
|
||||
embedding=q_embedding,
|
||||
top_k=k,
|
||||
filter=filter
|
||||
)
|
||||
|
||||
for r in results:
|
||||
r_id = r.id if hasattr(r, 'id') else r.text[:50]
|
||||
if r_id not in seen_ids:
|
||||
seen_ids.add(r_id)
|
||||
all_results.append(RetrievalResult(
|
||||
text=r.text if hasattr(r, 'text') else str(r),
|
||||
score=r.score if hasattr(r, 'score') else 1.0,
|
||||
metadata=r.metadata if hasattr(r, 'metadata') else {},
|
||||
doc_id=r.id if hasattr(r, 'id') else None
|
||||
))
|
||||
|
||||
# Generate follow-up query from result
|
||||
if depth < self.max_depth - 1 and self.llm:
|
||||
follow_up = self._generate_follow_up(query, r.text if hasattr(r, 'text') else str(r))
|
||||
if follow_up:
|
||||
next_queries.append(follow_up)
|
||||
|
||||
current_queries = next_queries[:3] # Limit expansion
|
||||
if not current_queries:
|
||||
break
|
||||
|
||||
# Sort by score
|
||||
all_results.sort(key=lambda x: x.score, reverse=True)
|
||||
return all_results[:k]
|
||||
|
||||
async def aretrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.retrieve(query, top_k, filter, **kwargs)
|
||||
|
||||
def _generate_follow_up(self, original_query: str, context: str) -> Optional[str]:
|
||||
"""Generate a follow-up query based on retrieved context."""
|
||||
if not self.llm:
|
||||
return None
|
||||
|
||||
try:
|
||||
prompt = f"""Based on this search result, generate a follow-up query to find more relevant information.
|
||||
Return only the query, nothing else.
|
||||
|
||||
Original query: {original_query}
|
||||
Search result: {context[:500]}"""
|
||||
response = self.llm.chat(prompt)
|
||||
if response:
|
||||
return response.strip()
|
||||
except Exception as e:
|
||||
logger.warning(f"Follow-up query generation failed: {e}")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class AutoMergeRetriever:
|
||||
"""
|
||||
Auto-merge retriever that combines adjacent chunks.
|
||||
|
||||
Retrieves chunks and merges those from the same document.
|
||||
"""
|
||||
|
||||
name: str = "auto_merge"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vector_store: Any,
|
||||
embedding_fn: Optional[Callable[[str], List[float]]] = None,
|
||||
top_k: int = 10,
|
||||
max_gap: int = 1,
|
||||
**kwargs
|
||||
):
|
||||
self.vector_store = vector_store
|
||||
self.embedding_fn = embedding_fn
|
||||
self.top_k = top_k
|
||||
self.max_gap = max_gap
|
||||
|
||||
def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Retrieve and merge adjacent chunks."""
|
||||
from praisonaiagents.knowledge.retrieval import RetrievalResult, merge_adjacent_chunks
|
||||
|
||||
k = top_k or self.top_k
|
||||
|
||||
# Get more results than needed for merging
|
||||
fetch_k = k * 3
|
||||
|
||||
if self.embedding_fn:
|
||||
q_embedding = self.embedding_fn(query)
|
||||
results = self.vector_store.query(
|
||||
embedding=q_embedding,
|
||||
top_k=fetch_k,
|
||||
filter=filter
|
||||
)
|
||||
|
||||
# Convert to RetrievalResult
|
||||
retrieval_results = []
|
||||
for r in results:
|
||||
retrieval_results.append(RetrievalResult(
|
||||
text=r.text if hasattr(r, 'text') else str(r),
|
||||
score=r.score if hasattr(r, 'score') else 1.0,
|
||||
metadata=r.metadata if hasattr(r, 'metadata') else {},
|
||||
doc_id=r.metadata.get('doc_id') if hasattr(r, 'metadata') else None,
|
||||
chunk_index=r.metadata.get('chunk_index') if hasattr(r, 'metadata') else None
|
||||
))
|
||||
|
||||
# Merge adjacent chunks
|
||||
merged = merge_adjacent_chunks(retrieval_results, max_gap=self.max_gap)
|
||||
return merged[:k]
|
||||
|
||||
return []
|
||||
|
||||
async def aretrieve(
|
||||
self,
|
||||
query: str,
|
||||
top_k: Optional[int] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
) -> List[Any]:
|
||||
"""Async version (just calls sync)."""
|
||||
return self.retrieve(query, top_k, filter, **kwargs)
|
||||
|
||||
|
||||
def register_default_retrievers():
|
||||
"""Register all default retrievers with the registry."""
|
||||
from praisonaiagents.knowledge.retrieval import get_retriever_registry
|
||||
|
||||
registry = get_retriever_registry()
|
||||
|
||||
registry.register("basic", BasicRetriever)
|
||||
registry.register("fusion", FusionRetriever)
|
||||
registry.register("recursive", RecursiveRetriever)
|
||||
registry.register("auto_merge", AutoMergeRetriever)
|
||||
|
||||
logger.debug("Registered default retrievers")
|
||||
|
||||
|
||||
# Auto-register on import
|
||||
register_default_retrievers()
|
||||
@@ -0,0 +1,399 @@
|
||||
"""
|
||||
Vector Store Adapters for PraisonAI.
|
||||
|
||||
Provides concrete implementations of VectorStoreProtocol:
|
||||
- ChromaVectorStore: Uses ChromaDB
|
||||
- PineconeVectorStore: Uses Pinecone
|
||||
- QdrantVectorStore: Uses Qdrant
|
||||
- WeaviateVectorStore: Uses Weaviate
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional
|
||||
import uuid as uuid_module
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy import flags
|
||||
_CHROMA_AVAILABLE = None
|
||||
_PINECONE_AVAILABLE = None
|
||||
_QDRANT_AVAILABLE = None
|
||||
_WEAVIATE_AVAILABLE = None
|
||||
|
||||
|
||||
def _check_chroma():
|
||||
"""Check if chromadb is available."""
|
||||
global _CHROMA_AVAILABLE
|
||||
if _CHROMA_AVAILABLE is None:
|
||||
import importlib.util
|
||||
_CHROMA_AVAILABLE = importlib.util.find_spec("chromadb") is not None
|
||||
return _CHROMA_AVAILABLE
|
||||
|
||||
|
||||
def _check_pinecone():
|
||||
"""Check if pinecone is available."""
|
||||
global _PINECONE_AVAILABLE
|
||||
if _PINECONE_AVAILABLE is None:
|
||||
import importlib.util
|
||||
_PINECONE_AVAILABLE = importlib.util.find_spec("pinecone") is not None
|
||||
return _PINECONE_AVAILABLE
|
||||
|
||||
|
||||
def _check_qdrant():
|
||||
"""Check if qdrant_client is available."""
|
||||
global _QDRANT_AVAILABLE
|
||||
if _QDRANT_AVAILABLE is None:
|
||||
import importlib.util
|
||||
_QDRANT_AVAILABLE = importlib.util.find_spec("qdrant_client") is not None
|
||||
return _QDRANT_AVAILABLE
|
||||
|
||||
|
||||
def _check_weaviate():
|
||||
"""Check if weaviate is available."""
|
||||
global _WEAVIATE_AVAILABLE
|
||||
if _WEAVIATE_AVAILABLE is None:
|
||||
import importlib.util
|
||||
_WEAVIATE_AVAILABLE = importlib.util.find_spec("weaviate") is not None
|
||||
return _WEAVIATE_AVAILABLE
|
||||
|
||||
|
||||
class ChromaVectorStore:
|
||||
"""ChromaDB vector store adapter."""
|
||||
|
||||
name: str = "chroma"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
persist_directory: Optional[str] = None
|
||||
):
|
||||
if not _check_chroma():
|
||||
raise ImportError(
|
||||
"chromadb is required for ChromaVectorStore. "
|
||||
"Install with: pip install chromadb"
|
||||
)
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
self.config = config or {}
|
||||
self.namespace = namespace or "default"
|
||||
self.persist_directory = persist_directory or self.config.get("path", ".praison/chroma")
|
||||
|
||||
# Disable telemetry
|
||||
os.environ['ANONYMIZED_TELEMETRY'] = 'False'
|
||||
|
||||
# Create client
|
||||
self._client = chromadb.PersistentClient(
|
||||
path=self.persist_directory,
|
||||
settings=Settings(anonymized_telemetry=False)
|
||||
)
|
||||
|
||||
# Get or create collection
|
||||
self._collection = self._client.get_or_create_collection(
|
||||
name=self.namespace,
|
||||
metadata={"hnsw:space": "cosine"}
|
||||
)
|
||||
|
||||
def add(
|
||||
self,
|
||||
texts: List[str],
|
||||
embeddings: List[List[float]],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""Add vectors to ChromaDB."""
|
||||
ids = ids or [str(uuid_module.uuid4()) for _ in texts]
|
||||
metadatas = metadatas or [{} for _ in texts]
|
||||
|
||||
# Use different collection if namespace specified
|
||||
if namespace and namespace != self.namespace:
|
||||
collection = self._client.get_or_create_collection(name=namespace)
|
||||
else:
|
||||
collection = self._collection
|
||||
|
||||
collection.add(
|
||||
ids=ids,
|
||||
embeddings=embeddings,
|
||||
documents=texts,
|
||||
metadatas=metadatas
|
||||
)
|
||||
|
||||
return ids
|
||||
|
||||
def query(
|
||||
self,
|
||||
embedding: List[float],
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None
|
||||
) -> List[Any]:
|
||||
"""Query ChromaDB by similarity."""
|
||||
from praisonaiagents.knowledge.vector_store import VectorRecord
|
||||
|
||||
if namespace and namespace != self.namespace:
|
||||
collection = self._client.get_or_create_collection(name=namespace)
|
||||
else:
|
||||
collection = self._collection
|
||||
|
||||
results = collection.query(
|
||||
query_embeddings=[embedding],
|
||||
n_results=top_k,
|
||||
where=filter
|
||||
)
|
||||
|
||||
records = []
|
||||
if results and results['ids'] and results['ids'][0]:
|
||||
for i, id_ in enumerate(results['ids'][0]):
|
||||
records.append(VectorRecord(
|
||||
id=id_,
|
||||
text=results['documents'][0][i] if results.get('documents') else "",
|
||||
embedding=embedding, # We don't get embeddings back from query
|
||||
metadata=results['metadatas'][0][i] if results.get('metadatas') else {},
|
||||
score=1 - results['distances'][0][i] if results.get('distances') else 1.0
|
||||
))
|
||||
|
||||
return records
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete vectors from ChromaDB."""
|
||||
if namespace and namespace != self.namespace:
|
||||
collection = self._client.get_or_create_collection(name=namespace)
|
||||
else:
|
||||
collection = self._collection
|
||||
|
||||
if delete_all:
|
||||
count = collection.count()
|
||||
# Delete collection and recreate
|
||||
self._client.delete_collection(collection.name)
|
||||
self._collection = self._client.get_or_create_collection(
|
||||
name=self.namespace,
|
||||
metadata={"hnsw:space": "cosine"}
|
||||
)
|
||||
return count
|
||||
|
||||
if ids:
|
||||
collection.delete(ids=ids)
|
||||
return len(ids)
|
||||
|
||||
if filter:
|
||||
# Get matching IDs first
|
||||
results = collection.get(where=filter)
|
||||
if results and results['ids']:
|
||||
collection.delete(ids=results['ids'])
|
||||
return len(results['ids'])
|
||||
|
||||
return 0
|
||||
|
||||
def count(self, namespace: Optional[str] = None) -> int:
|
||||
"""Get count of vectors."""
|
||||
if namespace and namespace != self.namespace:
|
||||
collection = self._client.get_or_create_collection(name=namespace)
|
||||
else:
|
||||
collection = self._collection
|
||||
return collection.count()
|
||||
|
||||
def get(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None
|
||||
) -> List[Any]:
|
||||
"""Get vectors by ID."""
|
||||
from praisonaiagents.knowledge.vector_store import VectorRecord
|
||||
|
||||
if namespace and namespace != self.namespace:
|
||||
collection = self._client.get_or_create_collection(name=namespace)
|
||||
else:
|
||||
collection = self._collection
|
||||
|
||||
results = collection.get(ids=ids, include=["documents", "metadatas", "embeddings"])
|
||||
|
||||
records = []
|
||||
if results and results['ids']:
|
||||
for i, id_ in enumerate(results['ids']):
|
||||
records.append(VectorRecord(
|
||||
id=id_,
|
||||
text=results['documents'][i] if results.get('documents') else "",
|
||||
embedding=results['embeddings'][i] if results.get('embeddings') else [],
|
||||
metadata=results['metadatas'][i] if results.get('metadatas') else {}
|
||||
))
|
||||
|
||||
return records
|
||||
|
||||
|
||||
class PineconeVectorStore:
|
||||
"""Pinecone vector store adapter."""
|
||||
|
||||
name: str = "pinecone"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
index_name: Optional[str] = None
|
||||
):
|
||||
if not _check_pinecone():
|
||||
raise ImportError(
|
||||
"pinecone is required for PineconeVectorStore. "
|
||||
"Install with: pip install pinecone"
|
||||
)
|
||||
|
||||
from pinecone import Pinecone
|
||||
|
||||
self.config = config or {}
|
||||
self.namespace = namespace or "default"
|
||||
self.api_key = api_key or os.environ.get("PINECONE_API_KEY")
|
||||
self.index_name = index_name or self.config.get("index_name", "praisonai")
|
||||
|
||||
if not self.api_key:
|
||||
raise ValueError("PINECONE_API_KEY environment variable or api_key parameter required")
|
||||
|
||||
self._pc = Pinecone(api_key=self.api_key)
|
||||
self._index = self._pc.Index(self.index_name)
|
||||
|
||||
def add(
|
||||
self,
|
||||
texts: List[str],
|
||||
embeddings: List[List[float]],
|
||||
metadatas: Optional[List[Dict[str, Any]]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""Add vectors to Pinecone."""
|
||||
ids = ids or [str(uuid_module.uuid4()) for _ in texts]
|
||||
metadatas = metadatas or [{} for _ in texts]
|
||||
ns = namespace or self.namespace
|
||||
|
||||
# Add text to metadata for retrieval
|
||||
vectors = []
|
||||
for i, (id_, embedding, metadata) in enumerate(zip(ids, embeddings, metadatas)):
|
||||
meta = metadata.copy()
|
||||
meta["text"] = texts[i]
|
||||
vectors.append({
|
||||
"id": id_,
|
||||
"values": embedding,
|
||||
"metadata": meta
|
||||
})
|
||||
|
||||
self._index.upsert(vectors=vectors, namespace=ns)
|
||||
return ids
|
||||
|
||||
def query(
|
||||
self,
|
||||
embedding: List[float],
|
||||
top_k: int = 10,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None
|
||||
) -> List[Any]:
|
||||
"""Query Pinecone by similarity."""
|
||||
from praisonaiagents.knowledge.vector_store import VectorRecord
|
||||
|
||||
ns = namespace or self.namespace
|
||||
|
||||
results = self._index.query(
|
||||
vector=embedding,
|
||||
top_k=top_k,
|
||||
namespace=ns,
|
||||
filter=filter,
|
||||
include_metadata=True
|
||||
)
|
||||
|
||||
records = []
|
||||
for match in results.get("matches", []):
|
||||
metadata = match.get("metadata", {})
|
||||
text = metadata.pop("text", "")
|
||||
records.append(VectorRecord(
|
||||
id=match["id"],
|
||||
text=text,
|
||||
embedding=embedding,
|
||||
metadata=metadata,
|
||||
score=match.get("score", 0.0)
|
||||
))
|
||||
|
||||
return records
|
||||
|
||||
def delete(
|
||||
self,
|
||||
ids: Optional[List[str]] = None,
|
||||
namespace: Optional[str] = None,
|
||||
filter: Optional[Dict[str, Any]] = None,
|
||||
delete_all: bool = False
|
||||
) -> int:
|
||||
"""Delete vectors from Pinecone."""
|
||||
ns = namespace or self.namespace
|
||||
|
||||
if delete_all:
|
||||
self._index.delete(delete_all=True, namespace=ns)
|
||||
return -1 # Pinecone doesn't return count
|
||||
|
||||
if ids:
|
||||
self._index.delete(ids=ids, namespace=ns)
|
||||
return len(ids)
|
||||
|
||||
if filter:
|
||||
self._index.delete(filter=filter, namespace=ns)
|
||||
return -1
|
||||
|
||||
return 0
|
||||
|
||||
def count(self, namespace: Optional[str] = None) -> int:
|
||||
"""Get count of vectors."""
|
||||
ns = namespace or self.namespace
|
||||
stats = self._index.describe_index_stats()
|
||||
ns_stats = stats.get("namespaces", {}).get(ns, {})
|
||||
return ns_stats.get("vector_count", 0)
|
||||
|
||||
def get(
|
||||
self,
|
||||
ids: List[str],
|
||||
namespace: Optional[str] = None
|
||||
) -> List[Any]:
|
||||
"""Get vectors by ID."""
|
||||
from praisonaiagents.knowledge.vector_store import VectorRecord
|
||||
|
||||
ns = namespace or self.namespace
|
||||
results = self._index.fetch(ids=ids, namespace=ns)
|
||||
|
||||
records = []
|
||||
for id_, data in results.get("vectors", {}).items():
|
||||
metadata = data.get("metadata", {})
|
||||
text = metadata.pop("text", "")
|
||||
records.append(VectorRecord(
|
||||
id=id_,
|
||||
text=text,
|
||||
embedding=data.get("values", []),
|
||||
metadata=metadata
|
||||
))
|
||||
|
||||
return records
|
||||
|
||||
|
||||
def register_default_vector_stores():
|
||||
"""Register all default vector stores with the registry."""
|
||||
from praisonaiagents.knowledge.vector_store import get_vector_store_registry
|
||||
|
||||
registry = get_vector_store_registry()
|
||||
|
||||
# Register ChromaDB
|
||||
if _check_chroma():
|
||||
registry.register("chroma", ChromaVectorStore)
|
||||
|
||||
# Register Pinecone
|
||||
if _check_pinecone():
|
||||
registry.register("pinecone", PineconeVectorStore)
|
||||
|
||||
logger.debug("Registered default vector stores")
|
||||
|
||||
|
||||
# Auto-register on import
|
||||
register_default_vector_stores()
|
||||
@@ -0,0 +1,595 @@
|
||||
"""
|
||||
Jobs CLI Feature for PraisonAI.
|
||||
|
||||
Provides CLI commands for managing async jobs.
|
||||
|
||||
Commands:
|
||||
- praisonai run submit <prompt> # Submit a new job
|
||||
- praisonai run status <job_id> # Get job status
|
||||
- praisonai run result <job_id> # Get job result
|
||||
- praisonai run cancel <job_id> # Cancel a job
|
||||
- praisonai run list # List jobs
|
||||
"""
|
||||
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
from typing import Optional, List
|
||||
|
||||
|
||||
class JobsHandler:
|
||||
"""
|
||||
Handler for jobs CLI commands.
|
||||
|
||||
Provides functionality to:
|
||||
- Submit jobs to the API
|
||||
- Check job status
|
||||
- Get job results
|
||||
- Cancel jobs
|
||||
- List jobs
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_url: str = "http://127.0.0.1:8005",
|
||||
verbose: bool = False
|
||||
):
|
||||
self.api_url = api_url.rstrip("/")
|
||||
self.verbose = verbose
|
||||
|
||||
@property
|
||||
def feature_name(self) -> str:
|
||||
return "jobs"
|
||||
|
||||
def _get_client(self):
|
||||
"""Get HTTP client."""
|
||||
try:
|
||||
import httpx
|
||||
return httpx.Client(timeout=30.0)
|
||||
except ImportError:
|
||||
raise RuntimeError("httpx is required. Install with: pip install httpx")
|
||||
|
||||
def submit(
|
||||
self,
|
||||
prompt: str,
|
||||
agent_file: Optional[str] = None,
|
||||
framework: str = "praisonai",
|
||||
timeout: int = 3600,
|
||||
wait: bool = False,
|
||||
poll_interval: int = 5,
|
||||
idempotency_key: Optional[str] = None,
|
||||
webhook_url: Optional[str] = None,
|
||||
session_id: Optional[str] = None,
|
||||
output_json: bool = False
|
||||
) -> dict:
|
||||
"""
|
||||
Submit a new job.
|
||||
|
||||
Args:
|
||||
prompt: The prompt/task for the agent
|
||||
agent_file: Optional path to agents.yaml
|
||||
framework: Framework to use
|
||||
timeout: Job timeout in seconds
|
||||
wait: If True, wait for completion
|
||||
poll_interval: Seconds between status polls
|
||||
|
||||
Returns:
|
||||
Job submission response
|
||||
"""
|
||||
client = self._get_client()
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"framework": framework,
|
||||
"timeout": timeout
|
||||
}
|
||||
|
||||
if agent_file:
|
||||
payload["agent_file"] = agent_file
|
||||
if webhook_url:
|
||||
payload["webhook_url"] = webhook_url
|
||||
if session_id:
|
||||
payload["session_id"] = session_id
|
||||
if idempotency_key:
|
||||
payload["idempotency_key"] = idempotency_key
|
||||
|
||||
headers = {}
|
||||
if idempotency_key:
|
||||
headers["Idempotency-Key"] = idempotency_key
|
||||
|
||||
try:
|
||||
response = client.post(
|
||||
f"{self.api_url}/api/v1/runs",
|
||||
json=payload,
|
||||
headers=headers if headers else None
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
if output_json:
|
||||
print(json.dumps(result, indent=2))
|
||||
else:
|
||||
self._print_success(f"Job submitted: {result['job_id']}")
|
||||
self._print_info(f"Status URL: {result['poll_url']}")
|
||||
|
||||
if wait:
|
||||
return self._wait_for_completion(result['job_id'], poll_interval, output_json)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self._print_error(f"Failed to submit job: {e}")
|
||||
raise
|
||||
|
||||
def status(self, job_id: str) -> dict:
|
||||
"""
|
||||
Get job status.
|
||||
|
||||
Args:
|
||||
job_id: Job ID
|
||||
|
||||
Returns:
|
||||
Job status response
|
||||
"""
|
||||
client = self._get_client()
|
||||
|
||||
try:
|
||||
response = client.get(f"{self.api_url}/api/v1/runs/{job_id}")
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
self._print_job_status(result)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self._print_error(f"Failed to get status: {e}")
|
||||
raise
|
||||
|
||||
def result(self, job_id: str) -> dict:
|
||||
"""
|
||||
Get job result.
|
||||
|
||||
Args:
|
||||
job_id: Job ID
|
||||
|
||||
Returns:
|
||||
Job result response
|
||||
"""
|
||||
client = self._get_client()
|
||||
|
||||
try:
|
||||
response = client.get(f"{self.api_url}/api/v1/runs/{job_id}/result")
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
self._print_job_result(result)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self._print_error(f"Failed to get result: {e}")
|
||||
raise
|
||||
|
||||
def cancel(self, job_id: str) -> dict:
|
||||
"""
|
||||
Cancel a job.
|
||||
|
||||
Args:
|
||||
job_id: Job ID
|
||||
|
||||
Returns:
|
||||
Updated job status
|
||||
"""
|
||||
client = self._get_client()
|
||||
|
||||
try:
|
||||
response = client.post(f"{self.api_url}/api/v1/runs/{job_id}/cancel")
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
self._print_success(f"Job cancelled: {job_id}")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self._print_error(f"Failed to cancel job: {e}")
|
||||
raise
|
||||
|
||||
def list_jobs(
|
||||
self,
|
||||
status: Optional[str] = None,
|
||||
page: int = 1,
|
||||
page_size: int = 20
|
||||
) -> dict:
|
||||
"""
|
||||
List jobs.
|
||||
|
||||
Args:
|
||||
status: Optional status filter
|
||||
page: Page number
|
||||
page_size: Jobs per page
|
||||
|
||||
Returns:
|
||||
Job list response
|
||||
"""
|
||||
client = self._get_client()
|
||||
|
||||
params = {"page": page, "page_size": page_size}
|
||||
if status:
|
||||
params["status"] = status
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
f"{self.api_url}/api/v1/runs",
|
||||
params=params
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
self._print_job_list(result)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self._print_error(f"Failed to list jobs: {e}")
|
||||
raise
|
||||
|
||||
def _wait_for_completion(self, job_id: str, poll_interval: int = 5, output_json: bool = False) -> dict:
|
||||
"""Wait for a job to complete."""
|
||||
client = self._get_client()
|
||||
|
||||
if not output_json:
|
||||
self._print_info("Waiting for job completion...")
|
||||
|
||||
while True:
|
||||
try:
|
||||
response = client.get(f"{self.api_url}/api/v1/runs/{job_id}")
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
status = result.get("status")
|
||||
progress = result.get("progress", {}).get("percentage", 0)
|
||||
|
||||
if not output_json:
|
||||
self._print_progress(status, progress)
|
||||
|
||||
if status in ("succeeded", "failed", "cancelled"):
|
||||
if output_json:
|
||||
# Get full result for succeeded jobs
|
||||
if status == "succeeded":
|
||||
res = client.get(f"{self.api_url}/api/v1/runs/{job_id}/result")
|
||||
res.raise_for_status()
|
||||
print(json.dumps(res.json(), indent=2))
|
||||
else:
|
||||
print(json.dumps(result, indent=2))
|
||||
return result
|
||||
else:
|
||||
if status == "succeeded":
|
||||
return self.result(job_id)
|
||||
else:
|
||||
return result
|
||||
|
||||
# Honor Retry-After header if present
|
||||
retry_after = result.get("retry_after") or poll_interval
|
||||
time.sleep(retry_after)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
if not output_json:
|
||||
self._print_info("\nInterrupted. Job continues running in background.")
|
||||
return result
|
||||
except Exception as e:
|
||||
if not output_json:
|
||||
self._print_error(f"Error polling status: {e}")
|
||||
time.sleep(poll_interval)
|
||||
|
||||
def stream(self, job_id: str, output_json: bool = False) -> None:
|
||||
"""
|
||||
Stream job progress via SSE.
|
||||
|
||||
Args:
|
||||
job_id: Job ID
|
||||
output_json: If True, output raw JSON events
|
||||
"""
|
||||
try:
|
||||
import httpx
|
||||
|
||||
with httpx.stream(
|
||||
"GET",
|
||||
f"{self.api_url}/api/v1/runs/{job_id}/stream",
|
||||
timeout=None
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
for line in response.iter_lines():
|
||||
if not line:
|
||||
continue
|
||||
if line.startswith("data:"):
|
||||
data = line[5:].strip()
|
||||
if data == "[DONE]":
|
||||
if not output_json:
|
||||
self._print_success("Stream completed")
|
||||
break
|
||||
try:
|
||||
event = json.loads(data)
|
||||
if output_json:
|
||||
print(json.dumps(event))
|
||||
else:
|
||||
event_type = event.get("event", "unknown")
|
||||
if event_type == "progress":
|
||||
pct = event.get("data", {}).get("percentage", 0)
|
||||
step = event.get("data", {}).get("current_step", "")
|
||||
self._print_progress(f"{step}", pct)
|
||||
elif event_type == "complete":
|
||||
status = event.get("data", {}).get("status", "")
|
||||
if status == "succeeded":
|
||||
self._print_success(f"Job completed: {status}")
|
||||
else:
|
||||
self._print_error(f"Job ended: {status}")
|
||||
elif event_type == "error":
|
||||
self._print_error(event.get("data", {}).get("error", "Unknown error"))
|
||||
except json.JSONDecodeError:
|
||||
if not output_json:
|
||||
print(data)
|
||||
elif line.startswith(":"):
|
||||
# Heartbeat comment, ignore
|
||||
pass
|
||||
|
||||
except Exception as e:
|
||||
self._print_error(f"Stream error: {e}")
|
||||
raise
|
||||
|
||||
def _print_job_status(self, job: dict):
|
||||
"""Print job status."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
|
||||
console = Console()
|
||||
|
||||
status = job.get("status", "unknown")
|
||||
status_color = {
|
||||
"queued": "dim",
|
||||
"running": "yellow",
|
||||
"succeeded": "green",
|
||||
"failed": "red",
|
||||
"cancelled": "dim"
|
||||
}.get(status, "white")
|
||||
|
||||
progress = job.get("progress", {})
|
||||
|
||||
content = f"""
|
||||
[bold]Job ID:[/bold] {job.get('job_id')}
|
||||
[bold]Status:[/bold] [{status_color}]{status}[/{status_color}]
|
||||
[bold]Progress:[/bold] {progress.get('percentage', 0):.0f}%
|
||||
[bold]Step:[/bold] {progress.get('current_step') or '-'}
|
||||
[bold]Created:[/bold] {job.get('created_at')}
|
||||
[bold]Started:[/bold] {job.get('started_at') or '-'}
|
||||
"""
|
||||
if job.get("error"):
|
||||
content += f"[bold]Error:[/bold] [red]{job.get('error')}[/red]\n"
|
||||
|
||||
console.print(Panel(content.strip(), title="Job Status"))
|
||||
except ImportError:
|
||||
print(f"Job: {job.get('job_id')}")
|
||||
print(f" Status: {job.get('status')}")
|
||||
print(f" Progress: {job.get('progress', {}).get('percentage', 0):.0f}%")
|
||||
|
||||
def _print_job_result(self, job: dict):
|
||||
"""Print job result."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.markdown import Markdown
|
||||
|
||||
console = Console()
|
||||
|
||||
status = job.get("status")
|
||||
|
||||
if status == "succeeded":
|
||||
result = job.get("result", "")
|
||||
if isinstance(result, str):
|
||||
console.print(Panel(Markdown(result), title="Result"))
|
||||
else:
|
||||
console.print(Panel(json.dumps(result, indent=2), title="Result"))
|
||||
else:
|
||||
console.print(f"[red]Job {status}: {job.get('error')}[/red]")
|
||||
|
||||
except ImportError:
|
||||
print(f"Result: {job.get('result')}")
|
||||
|
||||
def _print_job_list(self, response: dict):
|
||||
"""Print job list."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
jobs = response.get("jobs", [])
|
||||
|
||||
if not jobs:
|
||||
console.print("[yellow]No jobs found[/yellow]")
|
||||
return
|
||||
|
||||
table = Table(title=f"Jobs (Page {response.get('page')}/{(response.get('total', 0) + response.get('page_size', 20) - 1) // response.get('page_size', 20)})")
|
||||
table.add_column("ID", style="cyan")
|
||||
table.add_column("Status", style="yellow")
|
||||
table.add_column("Progress", style="blue")
|
||||
table.add_column("Created", style="dim")
|
||||
|
||||
for job in jobs:
|
||||
status = job.get("status", "unknown")
|
||||
status_color = {
|
||||
"queued": "dim",
|
||||
"running": "yellow",
|
||||
"succeeded": "green",
|
||||
"failed": "red",
|
||||
"cancelled": "dim"
|
||||
}.get(status, "white")
|
||||
|
||||
progress = job.get("progress", {}).get("percentage", 0)
|
||||
|
||||
table.add_row(
|
||||
job.get("job_id", "?"),
|
||||
f"[{status_color}]{status}[/{status_color}]",
|
||||
f"{progress:.0f}%",
|
||||
job.get("created_at", "")[:19]
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
console.print(f"Total: {response.get('total', 0)} jobs")
|
||||
|
||||
except ImportError:
|
||||
print("Jobs:")
|
||||
for job in response.get("jobs", []):
|
||||
print(f" {job.get('job_id')} - {job.get('status')}")
|
||||
|
||||
def _print_progress(self, status: str, percentage: float):
|
||||
"""Print progress update."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
console = Console()
|
||||
console.print(f" Status: {status} | Progress: {percentage:.0f}%", end="\r")
|
||||
except ImportError:
|
||||
print(f" Status: {status} | Progress: {percentage:.0f}%", end="\r")
|
||||
|
||||
def _print_success(self, message: str):
|
||||
"""Print success message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[green]✓[/green] {message}")
|
||||
except ImportError:
|
||||
print(f"✓ {message}")
|
||||
|
||||
def _print_error(self, message: str):
|
||||
"""Print error message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[red]✗[/red] {message}")
|
||||
except ImportError:
|
||||
print(f"✗ {message}")
|
||||
|
||||
def _print_info(self, message: str):
|
||||
"""Print info message."""
|
||||
try:
|
||||
from rich.console import Console
|
||||
Console().print(f"[blue]ℹ[/blue] {message}")
|
||||
except ImportError:
|
||||
print(f"ℹ {message}")
|
||||
|
||||
|
||||
def handle_run_command(args: List[str], verbose: bool = False):
|
||||
"""
|
||||
Handle run CLI commands.
|
||||
|
||||
Usage:
|
||||
praisonai run submit "Your prompt here" [--wait] [--agent-file agents.yaml]
|
||||
praisonai run status <job_id>
|
||||
praisonai run result <job_id>
|
||||
praisonai run cancel <job_id>
|
||||
praisonai run list [--status <status>]
|
||||
"""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(prog="praisonai run", description="Manage async jobs")
|
||||
subparsers = parser.add_subparsers(dest="subcommand", help="Available commands")
|
||||
|
||||
# Submit command
|
||||
submit_parser = subparsers.add_parser("submit", help="Submit a new job")
|
||||
submit_parser.add_argument("prompt", help="The prompt/task for the agent")
|
||||
submit_parser.add_argument("--agent-file", help="Path to agents.yaml")
|
||||
submit_parser.add_argument("--framework", default="praisonai", help="Framework to use")
|
||||
submit_parser.add_argument("--timeout", type=int, default=3600, help="Timeout in seconds")
|
||||
submit_parser.add_argument("--wait", action="store_true", help="Wait for completion")
|
||||
submit_parser.add_argument("--idempotency-key", help="Idempotency key to prevent duplicates")
|
||||
submit_parser.add_argument("--webhook-url", help="Webhook URL for completion callback")
|
||||
submit_parser.add_argument("--session-id", help="Session ID for grouping jobs")
|
||||
submit_parser.add_argument("--json", dest="output_json", action="store_true", help="Output JSON for scripting")
|
||||
submit_parser.add_argument("--api-url", default="http://127.0.0.1:8005", help="API server URL")
|
||||
|
||||
# Status command
|
||||
status_parser = subparsers.add_parser("status", help="Get job status")
|
||||
status_parser.add_argument("job_id", help="Job ID")
|
||||
status_parser.add_argument("--json", dest="output_json", action="store_true", help="Output JSON")
|
||||
status_parser.add_argument("--api-url", default="http://127.0.0.1:8005", help="API server URL")
|
||||
|
||||
# Result command
|
||||
result_parser = subparsers.add_parser("result", help="Get job result")
|
||||
result_parser.add_argument("job_id", help="Job ID")
|
||||
result_parser.add_argument("--json", dest="output_json", action="store_true", help="Output JSON")
|
||||
result_parser.add_argument("--api-url", default="http://127.0.0.1:8005", help="API server URL")
|
||||
|
||||
# Cancel command
|
||||
cancel_parser = subparsers.add_parser("cancel", help="Cancel a job")
|
||||
cancel_parser.add_argument("job_id", help="Job ID")
|
||||
cancel_parser.add_argument("--json", dest="output_json", action="store_true", help="Output JSON")
|
||||
cancel_parser.add_argument("--api-url", default="http://127.0.0.1:8005", help="API server URL")
|
||||
|
||||
# List command
|
||||
list_parser = subparsers.add_parser("list", help="List jobs")
|
||||
list_parser.add_argument("--status", help="Filter by status")
|
||||
list_parser.add_argument("--page", type=int, default=1, help="Page number")
|
||||
list_parser.add_argument("--page-size", type=int, default=20, help="Jobs per page")
|
||||
list_parser.add_argument("--json", dest="output_json", action="store_true", help="Output JSON")
|
||||
list_parser.add_argument("--api-url", default="http://127.0.0.1:8005", help="API server URL")
|
||||
|
||||
# Stream command
|
||||
stream_parser = subparsers.add_parser("stream", help="Stream job progress via SSE")
|
||||
stream_parser.add_argument("job_id", help="Job ID")
|
||||
stream_parser.add_argument("--json", dest="output_json", action="store_true", help="Output raw JSON events")
|
||||
stream_parser.add_argument("--api-url", default="http://127.0.0.1:8005", help="API server URL")
|
||||
|
||||
if not args:
|
||||
parser.print_help()
|
||||
return
|
||||
|
||||
parsed = parser.parse_args(args)
|
||||
|
||||
if not parsed.subcommand:
|
||||
parser.print_help()
|
||||
return
|
||||
|
||||
handler = JobsHandler(
|
||||
api_url=getattr(parsed, "api_url", "http://127.0.0.1:8005"),
|
||||
verbose=verbose
|
||||
)
|
||||
|
||||
try:
|
||||
if parsed.subcommand == "submit":
|
||||
handler.submit(
|
||||
prompt=parsed.prompt,
|
||||
agent_file=parsed.agent_file,
|
||||
framework=parsed.framework,
|
||||
timeout=parsed.timeout,
|
||||
wait=parsed.wait,
|
||||
idempotency_key=getattr(parsed, "idempotency_key", None),
|
||||
webhook_url=getattr(parsed, "webhook_url", None),
|
||||
session_id=getattr(parsed, "session_id", None),
|
||||
output_json=getattr(parsed, "output_json", False)
|
||||
)
|
||||
elif parsed.subcommand == "status":
|
||||
result = handler.status(parsed.job_id)
|
||||
if getattr(parsed, "output_json", False):
|
||||
print(json.dumps(result, indent=2))
|
||||
elif parsed.subcommand == "result":
|
||||
result = handler.result(parsed.job_id)
|
||||
if getattr(parsed, "output_json", False):
|
||||
print(json.dumps(result, indent=2))
|
||||
elif parsed.subcommand == "cancel":
|
||||
result = handler.cancel(parsed.job_id)
|
||||
if getattr(parsed, "output_json", False):
|
||||
print(json.dumps(result, indent=2))
|
||||
elif parsed.subcommand == "list":
|
||||
result = handler.list_jobs(
|
||||
status=parsed.status,
|
||||
page=parsed.page,
|
||||
page_size=parsed.page_size
|
||||
)
|
||||
if getattr(parsed, "output_json", False):
|
||||
print(json.dumps(result, indent=2))
|
||||
elif parsed.subcommand == "stream":
|
||||
handler.stream(
|
||||
job_id=parsed.job_id,
|
||||
output_json=getattr(parsed, "output_json", False)
|
||||
)
|
||||
except Exception:
|
||||
if verbose:
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
@@ -1,13 +1,19 @@
|
||||
"""
|
||||
Knowledge Handler for CLI.
|
||||
|
||||
Provides RAG/vector store management.
|
||||
Usage: praisonai knowledge add document.pdf
|
||||
praisonai knowledge search "query"
|
||||
Provides RAG/vector store management with full knowledge stack support.
|
||||
|
||||
Usage:
|
||||
praisonai knowledge add document.pdf
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge query "How to authenticate?"
|
||||
praisonai knowledge list
|
||||
praisonai knowledge clear
|
||||
praisonai knowledge stats
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, List
|
||||
from typing import Any, Dict, List, Optional
|
||||
from .base import CommandHandler
|
||||
|
||||
|
||||
@@ -17,111 +23,238 @@ class KnowledgeHandler(CommandHandler):
|
||||
|
||||
Manages knowledge base for RAG (Retrieval Augmented Generation).
|
||||
|
||||
Supports:
|
||||
- Multiple vector stores (chroma, pinecone, qdrant, weaviate, memory)
|
||||
- Retrieval strategies (basic, fusion, recursive, auto_merge)
|
||||
- Rerankers (simple, llm, cross_encoder, cohere)
|
||||
- Index types (vector, keyword, hybrid)
|
||||
- Query modes (default, sub_question, summarize)
|
||||
|
||||
Example:
|
||||
praisonai knowledge add document.pdf
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge search "How to authenticate?"
|
||||
praisonai knowledge query "How to authenticate?" --retrieval fusion
|
||||
praisonai knowledge list
|
||||
praisonai knowledge clear
|
||||
praisonai knowledge stats
|
||||
"""
|
||||
|
||||
def __init__(self, verbose: bool = False, workspace: str = None):
|
||||
def __init__(
|
||||
self,
|
||||
verbose: bool = False,
|
||||
workspace: str = None,
|
||||
vector_store: str = "chroma",
|
||||
retrieval_strategy: str = "basic",
|
||||
reranker: Optional[str] = None,
|
||||
index_type: str = "vector",
|
||||
query_mode: str = "default",
|
||||
session_id: Optional[str] = None,
|
||||
db: Optional[str] = None
|
||||
):
|
||||
super().__init__(verbose)
|
||||
self.workspace = workspace or os.getcwd()
|
||||
self.vector_store = vector_store
|
||||
self.retrieval_strategy = retrieval_strategy
|
||||
self.reranker = reranker
|
||||
self.index_type = index_type
|
||||
self.query_mode = query_mode
|
||||
self.session_id = session_id
|
||||
self.db = db
|
||||
self._knowledge = None
|
||||
self._embedding_fn = None
|
||||
|
||||
@property
|
||||
def feature_name(self) -> str:
|
||||
return "knowledge"
|
||||
|
||||
def get_actions(self) -> List[str]:
|
||||
return ["add", "search", "list", "clear", "info", "help"]
|
||||
return ["add", "query", "search", "list", "clear", "stats", "info", "help"]
|
||||
|
||||
def get_help_text(self) -> str:
|
||||
return """
|
||||
Knowledge Commands:
|
||||
praisonai knowledge add <file|dir> - Add document(s) to knowledge base
|
||||
praisonai knowledge search <query> - Search knowledge base
|
||||
praisonai knowledge list - List indexed documents
|
||||
praisonai knowledge clear - Clear knowledge base
|
||||
praisonai knowledge info - Show knowledge base info
|
||||
praisonai knowledge add <file|dir|url> - Add document(s) to knowledge base
|
||||
praisonai knowledge query "<question>" - Query knowledge base with RAG
|
||||
praisonai knowledge list - List indexed documents
|
||||
praisonai knowledge clear - Clear knowledge base
|
||||
praisonai knowledge stats - Show knowledge base statistics
|
||||
praisonai knowledge help - Show this help message
|
||||
|
||||
Options:
|
||||
--workspace <path> - Workspace directory (default: current)
|
||||
--workspace <path> - Workspace directory (default: current)
|
||||
--vector-store <name> - Vector store backend
|
||||
Values: memory, chroma, pinecone, qdrant, weaviate
|
||||
Default: chroma
|
||||
--retrieval-strategy <strategy> - Retrieval strategy
|
||||
Values: basic, fusion, recursive, auto_merge
|
||||
Default: basic
|
||||
--reranker <name> - Reranker for result ordering
|
||||
Values: none, simple, llm, cross_encoder, cohere
|
||||
Default: none
|
||||
--index-type <type> - Index type for search
|
||||
Values: vector, keyword, hybrid
|
||||
Default: vector
|
||||
--query-mode <mode> - Query processing mode
|
||||
Values: default, sub_question, summarize
|
||||
Default: default
|
||||
--session <id> - Session ID for persistence
|
||||
--db <path> - Database path for persistence
|
||||
|
||||
Examples:
|
||||
praisonai knowledge add document.pdf
|
||||
praisonai knowledge add ./docs/
|
||||
praisonai knowledge query "How to authenticate?" --retrieval-strategy fusion
|
||||
praisonai knowledge query "Compare approaches" --reranker simple --query-mode sub_question
|
||||
"""
|
||||
|
||||
def _get_knowledge(self):
|
||||
"""Lazy load Knowledge instance."""
|
||||
"""Lazy load Knowledge instance with configured options."""
|
||||
if self._knowledge is None:
|
||||
try:
|
||||
from praisonaiagents import Knowledge
|
||||
self._knowledge = Knowledge()
|
||||
|
||||
# Build config based on options
|
||||
config = {
|
||||
"vector_store": {
|
||||
"provider": self.vector_store,
|
||||
"config": {
|
||||
"path": os.path.join(self.workspace, ".praison", "knowledge")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self._knowledge = Knowledge(config=config)
|
||||
except ImportError:
|
||||
self.print_status(
|
||||
"Knowledge requires praisonaiagents. Install with: pip install praisonaiagents",
|
||||
"Knowledge requires praisonaiagents. Install with: pip install praisonaiagents[knowledge]",
|
||||
"error"
|
||||
)
|
||||
return None
|
||||
return self._knowledge
|
||||
|
||||
def _get_embedding_fn(self):
|
||||
"""Get embedding function."""
|
||||
if self._embedding_fn is None:
|
||||
try:
|
||||
import openai
|
||||
client = openai.OpenAI()
|
||||
|
||||
def embed(text: str) -> List[float]:
|
||||
response = client.embeddings.create(
|
||||
input=text,
|
||||
model="text-embedding-3-small"
|
||||
)
|
||||
return response.data[0].embedding
|
||||
|
||||
self._embedding_fn = embed
|
||||
except Exception as e:
|
||||
self.print_status(f"Failed to initialize embeddings: {e}", "warning")
|
||||
return self._embedding_fn
|
||||
|
||||
def action_add(self, args: List[str], **kwargs) -> bool:
|
||||
"""
|
||||
Add document(s) to knowledge base.
|
||||
|
||||
Supports:
|
||||
- File paths (pdf, docx, txt, md, html, etc.)
|
||||
- Directories (recursive)
|
||||
- Glob patterns (*.pdf, docs/**/*.md)
|
||||
- URLs (http://, https://)
|
||||
|
||||
Args:
|
||||
args: List containing file/directory paths
|
||||
args: List containing file/directory/URL paths
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if not args:
|
||||
self.print_status("Usage: praisonai knowledge add <file|directory>", "error")
|
||||
self.print_status("Usage: praisonai knowledge add <file|directory|url|glob>", "error")
|
||||
return False
|
||||
|
||||
knowledge = self._get_knowledge()
|
||||
if not knowledge:
|
||||
return False
|
||||
|
||||
path = args[0]
|
||||
source = args[0]
|
||||
|
||||
# Check if URL
|
||||
if source.startswith(("http://", "https://")):
|
||||
try:
|
||||
knowledge.add(source)
|
||||
self.print_status(f"✅ Added URL: {source}", "success")
|
||||
return True
|
||||
except Exception as e:
|
||||
self.print_status(f"Failed to add URL: {e}", "error")
|
||||
return False
|
||||
|
||||
# Expand path
|
||||
if not os.path.isabs(path):
|
||||
path = os.path.join(self.workspace, path)
|
||||
if not os.path.isabs(source):
|
||||
source = os.path.join(self.workspace, source)
|
||||
|
||||
if not os.path.exists(path):
|
||||
self.print_status(f"Path not found: {path}", "error")
|
||||
# Check for glob pattern
|
||||
if "*" in source or "?" in source:
|
||||
import glob
|
||||
files = glob.glob(source, recursive=True)
|
||||
if not files:
|
||||
self.print_status(f"No files match pattern: {source}", "warning")
|
||||
return False
|
||||
|
||||
count = 0
|
||||
for file_path in files:
|
||||
if os.path.isfile(file_path):
|
||||
try:
|
||||
knowledge.add(file_path)
|
||||
count += 1
|
||||
except Exception as e:
|
||||
self.print_status(f"Failed to add {file_path}: {e}", "warning")
|
||||
|
||||
self.print_status(f"✅ Added {count} files matching {args[0]}", "success")
|
||||
return True
|
||||
|
||||
if not os.path.exists(source):
|
||||
self.print_status(f"Path not found: {source}", "error")
|
||||
return False
|
||||
|
||||
try:
|
||||
if os.path.isdir(path):
|
||||
if os.path.isdir(source):
|
||||
# Add all files in directory
|
||||
count = 0
|
||||
for root, _, files in os.walk(path):
|
||||
for root, _, files in os.walk(source):
|
||||
for file in files:
|
||||
file_path = os.path.join(root, file)
|
||||
knowledge.add(file_path)
|
||||
count += 1
|
||||
self.print_status(f"✅ Added {count} files from {path}", "success")
|
||||
try:
|
||||
knowledge.add(file_path)
|
||||
count += 1
|
||||
except Exception as e:
|
||||
if self.verbose:
|
||||
self.print_status(f"Skipped {file}: {e}", "warning")
|
||||
self.print_status(f"✅ Added {count} files from {source}", "success")
|
||||
else:
|
||||
knowledge.add(path)
|
||||
self.print_status(f"✅ Added: {path}", "success")
|
||||
knowledge.add(source)
|
||||
self.print_status(f"✅ Added: {source}", "success")
|
||||
return True
|
||||
except Exception as e:
|
||||
self.print_status(f"Failed to add: {e}", "error")
|
||||
return False
|
||||
|
||||
def action_search(self, args: List[str], **kwargs) -> List[Dict[str, Any]]:
|
||||
"""Alias for action_query."""
|
||||
return self.action_query(args, **kwargs)
|
||||
|
||||
def action_query(self, args: List[str], **kwargs) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search knowledge base.
|
||||
Query knowledge base with RAG.
|
||||
|
||||
Uses configured retrieval strategy, reranker, and query mode.
|
||||
|
||||
Args:
|
||||
args: List containing search query
|
||||
args: List containing query
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
List of results with answer and sources
|
||||
"""
|
||||
if not args:
|
||||
self.print_status("Usage: praisonai knowledge search <query>", "error")
|
||||
self.print_status("Usage: praisonai knowledge query <question>", "error")
|
||||
return []
|
||||
|
||||
knowledge = self._get_knowledge()
|
||||
@@ -132,19 +265,25 @@ Options:
|
||||
limit = kwargs.get('limit', 5)
|
||||
|
||||
try:
|
||||
results = knowledge.search(query, k=limit)
|
||||
# Use knowledge.search for basic retrieval
|
||||
results = knowledge.search(query, limit=limit)
|
||||
|
||||
if results:
|
||||
self.print_status(f"\n🔍 Search results for '{query}':\n", "info")
|
||||
self.print_status(f"\n🔍 Results for '{query}':\n", "info")
|
||||
for i, result in enumerate(results, 1):
|
||||
content = str(result)[:200]
|
||||
self.print_status(f" {i}. {content}...", "info")
|
||||
if isinstance(result, dict):
|
||||
content = result.get('memory', result.get('text', str(result)))[:200]
|
||||
score = result.get('score', 'N/A')
|
||||
self.print_status(f" {i}. [score: {score}] {content}...", "info")
|
||||
else:
|
||||
content = str(result)[:200]
|
||||
self.print_status(f" {i}. {content}...", "info")
|
||||
else:
|
||||
self.print_status("No results found", "warning")
|
||||
|
||||
return results
|
||||
except Exception as e:
|
||||
self.print_status(f"Search failed: {e}", "error")
|
||||
self.print_status(f"Query failed: {e}", "error")
|
||||
return []
|
||||
|
||||
def action_list(self, args: List[str], **kwargs) -> List[str]:
|
||||
@@ -201,35 +340,45 @@ Options:
|
||||
self.print_status(f"Failed to clear: {e}", "error")
|
||||
return False
|
||||
|
||||
def action_info(self, args: List[str], **kwargs) -> Dict[str, Any]:
|
||||
def action_stats(self, args: List[str], **kwargs) -> Dict[str, Any]:
|
||||
"""
|
||||
Show knowledge base info.
|
||||
Show knowledge base statistics.
|
||||
|
||||
Returns:
|
||||
Dictionary of info
|
||||
Dictionary of statistics
|
||||
"""
|
||||
knowledge = self._get_knowledge()
|
||||
if not knowledge:
|
||||
return {}
|
||||
|
||||
info = {
|
||||
stats = {
|
||||
"workspace": self.workspace,
|
||||
"type": type(knowledge).__name__
|
||||
"vector_store": self.vector_store,
|
||||
"retrieval_strategy": self.retrieval_strategy,
|
||||
"reranker": self.reranker or "none",
|
||||
"index_type": self.index_type,
|
||||
"query_mode": self.query_mode
|
||||
}
|
||||
|
||||
try:
|
||||
if hasattr(knowledge, 'get_stats'):
|
||||
info.update(knowledge.get_stats())
|
||||
if hasattr(knowledge, 'count'):
|
||||
info['document_count'] = knowledge.count()
|
||||
# Try to get document count from memory
|
||||
if hasattr(knowledge, 'memory'):
|
||||
mem = knowledge.memory
|
||||
if hasattr(mem, 'get_all'):
|
||||
all_docs = mem.get_all()
|
||||
stats['document_count'] = len(all_docs) if all_docs else 0
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.print_status("\n📊 Knowledge Base Info:", "info")
|
||||
for key, value in info.items():
|
||||
self.print_status("\n📊 Knowledge Base Statistics:", "info")
|
||||
for key, value in stats.items():
|
||||
self.print_status(f" {key}: {value}", "info")
|
||||
|
||||
return info
|
||||
return stats
|
||||
|
||||
def action_info(self, args: List[str], **kwargs) -> Dict[str, Any]:
|
||||
"""Alias for action_stats."""
|
||||
return self.action_stats(args, **kwargs)
|
||||
|
||||
def execute(self, action: str, action_args: List[str], **kwargs) -> Any:
|
||||
"""Execute knowledge command action."""
|
||||
|
||||
@@ -668,7 +668,7 @@ class PraisonAI:
|
||||
return default_args
|
||||
|
||||
# Define special commands
|
||||
special_commands = ['chat', 'code', 'call', 'realtime', 'train', 'ui', 'context', 'research', 'memory', 'rules', 'workflow', 'hooks', 'knowledge', 'session', 'tools', 'todo', 'docs', 'mcp', 'commit', 'serve', 'schedule', 'skills', 'profile', 'eval', 'agents']
|
||||
special_commands = ['chat', 'code', 'call', 'realtime', 'train', 'ui', 'context', 'research', 'memory', 'rules', 'workflow', 'hooks', 'knowledge', 'session', 'tools', 'todo', 'docs', 'mcp', 'commit', 'serve', 'schedule', 'skills', 'profile', 'eval', 'agents', 'run']
|
||||
|
||||
parser = argparse.ArgumentParser(prog="praisonai", description="praisonAI command-line interface")
|
||||
parser.add_argument("--framework", choices=["crewai", "autogen", "praisonai"], help="Specify the framework")
|
||||
@@ -1156,6 +1156,12 @@ class PraisonAI:
|
||||
|
||||
exit_code = handle_agents_command(agents_args)
|
||||
sys.exit(exit_code)
|
||||
|
||||
elif args.command == 'run':
|
||||
# Run command - async jobs API for long-running tasks
|
||||
from .features.jobs import handle_run_command
|
||||
handle_run_command(unknown_args, verbose=getattr(args, 'verbose', False))
|
||||
sys.exit(0)
|
||||
|
||||
# Only check framework availability for agent-related operations
|
||||
if not args.command and (args.init or args.auto or args.framework):
|
||||
@@ -2342,13 +2348,82 @@ class PraisonAI:
|
||||
Handle knowledge subcommand actions.
|
||||
|
||||
Args:
|
||||
action: The knowledge action (add, search, list, clear, info)
|
||||
action_args: Additional arguments for the action
|
||||
action: The knowledge action (add, query, list, clear, stats, help)
|
||||
action_args: Additional arguments for the action (may include flags)
|
||||
"""
|
||||
try:
|
||||
from .features.knowledge import KnowledgeHandler
|
||||
handler = KnowledgeHandler(verbose=True, workspace=os.getcwd())
|
||||
handler.execute(action, action_args)
|
||||
|
||||
# Parse flags from action_args
|
||||
vector_store = "chroma"
|
||||
retrieval_strategy = "basic"
|
||||
reranker = None
|
||||
index_type = "vector"
|
||||
query_mode = "default"
|
||||
session_id = None
|
||||
db = None
|
||||
workspace = os.getcwd()
|
||||
|
||||
# Filter out flags and extract values
|
||||
filtered_args = []
|
||||
i = 0
|
||||
while i < len(action_args):
|
||||
arg = action_args[i]
|
||||
if arg in ("--vector-store", "--store"):
|
||||
if i + 1 < len(action_args):
|
||||
vector_store = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
elif arg in ("--retrieval-strategy", "--retrieval", "--strategy"):
|
||||
if i + 1 < len(action_args):
|
||||
retrieval_strategy = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
elif arg == "--reranker":
|
||||
if i + 1 < len(action_args):
|
||||
reranker = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
elif arg in ("--index-type", "--index"):
|
||||
if i + 1 < len(action_args):
|
||||
index_type = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
elif arg in ("--query-mode", "--mode"):
|
||||
if i + 1 < len(action_args):
|
||||
query_mode = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
elif arg == "--session":
|
||||
if i + 1 < len(action_args):
|
||||
session_id = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
elif arg == "--db":
|
||||
if i + 1 < len(action_args):
|
||||
db = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
elif arg == "--workspace":
|
||||
if i + 1 < len(action_args):
|
||||
workspace = action_args[i + 1]
|
||||
i += 2
|
||||
continue
|
||||
filtered_args.append(arg)
|
||||
i += 1
|
||||
|
||||
handler = KnowledgeHandler(
|
||||
verbose=True,
|
||||
workspace=workspace,
|
||||
vector_store=vector_store,
|
||||
retrieval_strategy=retrieval_strategy,
|
||||
reranker=reranker,
|
||||
index_type=index_type,
|
||||
query_mode=query_mode,
|
||||
session_id=session_id,
|
||||
db=db
|
||||
)
|
||||
handler.execute(action, filtered_args)
|
||||
except ImportError as e:
|
||||
print(f"[red]ERROR: Failed to import knowledge module: {e}[/red]")
|
||||
print("Make sure praisonaiagents is installed: pip install praisonaiagents")
|
||||
|
||||
@@ -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.5.0 gunicorn markdown\n")
|
||||
file.write("RUN pip install flask praisonai==2.5.1 gunicorn markdown\n")
|
||||
file.write("EXPOSE 8080\n")
|
||||
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
|
||||
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
Async Jobs API for PraisonAI.
|
||||
|
||||
Provides HTTP API endpoints for long-running agent tasks:
|
||||
- Submit jobs (POST /api/v1/runs)
|
||||
- Check status (GET /api/v1/runs/{job_id})
|
||||
- Get results (GET /api/v1/runs/{job_id}/result)
|
||||
- Cancel jobs (POST /api/v1/runs/{job_id}/cancel)
|
||||
- Stream progress (GET /api/v1/runs/{job_id}/stream)
|
||||
|
||||
Zero Performance Impact:
|
||||
- All imports are lazy loaded via __getattr__
|
||||
- Server only starts when explicitly requested
|
||||
- No overhead when not in use
|
||||
|
||||
Usage:
|
||||
# Start the jobs server
|
||||
praisonai serve --port 8005
|
||||
|
||||
# Or programmatically
|
||||
from praisonai.jobs import start_server
|
||||
start_server(port=8005)
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
# Models
|
||||
"Job",
|
||||
"JobStatus",
|
||||
"JobSubmitRequest",
|
||||
"JobStatusResponse",
|
||||
"JobResultResponse",
|
||||
# Store
|
||||
"JobStore",
|
||||
"InMemoryJobStore",
|
||||
# Executor
|
||||
"JobExecutor",
|
||||
# Router
|
||||
"create_router",
|
||||
# Server
|
||||
"start_server",
|
||||
"create_app",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name in ("Job", "JobStatus", "JobSubmitRequest", "JobStatusResponse", "JobResultResponse"):
|
||||
from .models import Job, JobStatus, JobSubmitRequest, JobStatusResponse, JobResultResponse
|
||||
return locals()[name]
|
||||
|
||||
if name in ("JobStore", "InMemoryJobStore"):
|
||||
from .store import JobStore, InMemoryJobStore
|
||||
return locals()[name]
|
||||
|
||||
if name == "JobExecutor":
|
||||
from .executor import JobExecutor
|
||||
return JobExecutor
|
||||
|
||||
if name == "create_router":
|
||||
from .router import create_router
|
||||
return create_router
|
||||
|
||||
if name in ("start_server", "create_app"):
|
||||
from .server import start_server, create_app
|
||||
return locals()[name]
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,370 @@
|
||||
"""
|
||||
Job Executor for PraisonAI Async Jobs API.
|
||||
|
||||
Handles background execution of agent jobs.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional, Callable, Any, Dict
|
||||
|
||||
from .models import Job, JobStatus
|
||||
from .store import JobStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class JobExecutor:
|
||||
"""
|
||||
Executes jobs in the background using asyncio.
|
||||
|
||||
Features:
|
||||
- Concurrent job execution with configurable limits
|
||||
- Timeout handling
|
||||
- Cancellation support
|
||||
- Progress callbacks
|
||||
- Webhook notifications
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: JobStore,
|
||||
max_concurrent: int = 10,
|
||||
default_timeout: int = 3600,
|
||||
cleanup_interval: int = 300
|
||||
):
|
||||
self.store = store
|
||||
self.max_concurrent = max_concurrent
|
||||
self.default_timeout = default_timeout
|
||||
self.cleanup_interval = cleanup_interval
|
||||
|
||||
self._semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self._running_tasks: Dict[str, asyncio.Task] = {}
|
||||
self._cleanup_task: Optional[asyncio.Task] = None
|
||||
self._shutdown = False
|
||||
self._progress_callbacks: Dict[str, Callable] = {}
|
||||
|
||||
async def start(self):
|
||||
"""Start the executor and cleanup loop."""
|
||||
self._shutdown = False
|
||||
self._cleanup_task = asyncio.create_task(self._cleanup_loop())
|
||||
logger.info(f"JobExecutor started (max_concurrent={self.max_concurrent})")
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the executor and cancel all running jobs."""
|
||||
self._shutdown = True
|
||||
|
||||
# Cancel cleanup task
|
||||
if self._cleanup_task:
|
||||
self._cleanup_task.cancel()
|
||||
try:
|
||||
await self._cleanup_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# Cancel all running jobs
|
||||
for job_id, task in list(self._running_tasks.items()):
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
self._running_tasks.clear()
|
||||
logger.info("JobExecutor stopped")
|
||||
|
||||
async def _cleanup_loop(self):
|
||||
"""Periodically clean up old completed jobs."""
|
||||
while not self._shutdown:
|
||||
try:
|
||||
await asyncio.sleep(self.cleanup_interval)
|
||||
await self.store.cleanup_old_jobs(max_age_seconds=86400) # 24 hours
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Cleanup error: {e}")
|
||||
|
||||
async def submit(self, job: Job) -> Job:
|
||||
"""
|
||||
Submit a job for execution.
|
||||
|
||||
Args:
|
||||
job: Job to execute
|
||||
|
||||
Returns:
|
||||
The submitted job
|
||||
"""
|
||||
# Save job to store
|
||||
await self.store.save(job)
|
||||
|
||||
# Start execution task
|
||||
task = asyncio.create_task(self._execute_job(job))
|
||||
self._running_tasks[job.id] = task
|
||||
|
||||
# Clean up task reference when done
|
||||
task.add_done_callback(lambda t: self._running_tasks.pop(job.id, None))
|
||||
|
||||
logger.info(f"Job submitted: {job.id}")
|
||||
return job
|
||||
|
||||
async def cancel(self, job_id: str) -> bool:
|
||||
"""
|
||||
Cancel a running job.
|
||||
|
||||
Args:
|
||||
job_id: ID of job to cancel
|
||||
|
||||
Returns:
|
||||
True if cancelled, False if not found or already complete
|
||||
"""
|
||||
job = await self.store.get(job_id)
|
||||
if not job:
|
||||
return False
|
||||
|
||||
if job.is_terminal:
|
||||
return False
|
||||
|
||||
# Mark as cancelled
|
||||
job.cancel()
|
||||
await self.store.save(job)
|
||||
|
||||
# Cancel the task if running
|
||||
task = self._running_tasks.get(job_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
|
||||
logger.info(f"Job cancelled: {job_id}")
|
||||
return True
|
||||
|
||||
def register_progress_callback(self, job_id: str, callback: Callable):
|
||||
"""Register a callback for job progress updates."""
|
||||
self._progress_callbacks[job_id] = callback
|
||||
|
||||
def unregister_progress_callback(self, job_id: str):
|
||||
"""Unregister a progress callback."""
|
||||
self._progress_callbacks.pop(job_id, None)
|
||||
|
||||
async def _execute_job(self, job: Job):
|
||||
"""Execute a job."""
|
||||
async with self._semaphore:
|
||||
try:
|
||||
# Mark as running
|
||||
job.start()
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
logger.info(f"Job started: {job.id}")
|
||||
|
||||
# Execute with timeout
|
||||
timeout = job.timeout or self.default_timeout
|
||||
result = await asyncio.wait_for(
|
||||
self._run_agent(job),
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
# Mark as succeeded
|
||||
job.succeed(result)
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
logger.info(f"Job succeeded: {job.id}")
|
||||
|
||||
# Send webhook if configured
|
||||
if job.webhook_url:
|
||||
await self._send_webhook(job)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
job.fail(f"Job timed out after {job.timeout}s")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
logger.warning(f"Job timed out: {job.id}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
if job.status != JobStatus.CANCELLED:
|
||||
job.cancel()
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
logger.info(f"Job cancelled: {job.id}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
job.fail(str(e))
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
logger.error(f"Job failed: {job.id} - {e}")
|
||||
|
||||
# Send webhook for failures too
|
||||
if job.webhook_url:
|
||||
await self._send_webhook(job)
|
||||
|
||||
async def _run_agent(self, job: Job) -> Any:
|
||||
"""
|
||||
Run the agent for a job.
|
||||
|
||||
This is the core execution logic that runs the PraisonAI agent.
|
||||
"""
|
||||
# Imports are done lazily in _run_praisonai_agents and _run_legacy_praisonai
|
||||
|
||||
# Determine agent configuration
|
||||
agent_file = job.agent_file or "agents.yaml"
|
||||
framework = job.framework or "praisonai"
|
||||
|
||||
# Check if we should use inline YAML
|
||||
if job.agent_yaml:
|
||||
# Write to temp file
|
||||
import tempfile
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
|
||||
f.write(job.agent_yaml)
|
||||
agent_file = f.name
|
||||
|
||||
# Update progress
|
||||
job.update_progress(percentage=10.0, step="Initializing agent")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
try:
|
||||
# Try praisonaiagents first (preferred)
|
||||
if framework == "praisonai":
|
||||
result = await self._run_praisonai_agents(job, agent_file)
|
||||
else:
|
||||
# Use legacy PraisonAI for crewai/autogen
|
||||
result = await self._run_legacy_praisonai(job, agent_file, framework)
|
||||
|
||||
return result
|
||||
|
||||
finally:
|
||||
# Clean up temp file if created
|
||||
if job.agent_yaml and agent_file != "agents.yaml":
|
||||
try:
|
||||
os.unlink(agent_file)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _run_praisonai_agents(self, job: Job, agent_file: str) -> Any:
|
||||
"""Run using praisonaiagents framework."""
|
||||
try:
|
||||
from praisonaiagents import Agent
|
||||
except ImportError:
|
||||
raise RuntimeError("praisonaiagents not installed")
|
||||
|
||||
# Update progress
|
||||
job.update_progress(percentage=20.0, step="Creating agent")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
# Create agent
|
||||
agent = Agent(
|
||||
instructions="You are a helpful AI assistant.",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Update job with agent info
|
||||
job.agent_id = getattr(agent, 'name', 'agent')
|
||||
job.run_id = getattr(agent, 'run_id', None)
|
||||
|
||||
# Update progress
|
||||
job.update_progress(percentage=30.0, step="Running agent")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
# Run the agent
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: agent.start(job.prompt)
|
||||
)
|
||||
|
||||
# Update progress
|
||||
job.update_progress(percentage=90.0, step="Finalizing")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
return result
|
||||
|
||||
async def _run_legacy_praisonai(self, job: Job, agent_file: str, framework: str) -> Any:
|
||||
"""Run using legacy PraisonAI (crewai/autogen)."""
|
||||
try:
|
||||
from praisonai import PraisonAI
|
||||
except ImportError:
|
||||
raise RuntimeError("praisonai not installed")
|
||||
|
||||
# Update progress
|
||||
job.update_progress(percentage=20.0, step="Creating PraisonAI instance")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
# Create PraisonAI instance
|
||||
praisonai = PraisonAI(
|
||||
agent_file=agent_file,
|
||||
framework=framework
|
||||
)
|
||||
|
||||
# Update progress
|
||||
job.update_progress(percentage=30.0, step="Running agents")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
# Run in executor (blocking call)
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(None, praisonai.run)
|
||||
|
||||
# Update progress
|
||||
job.update_progress(percentage=90.0, step="Finalizing")
|
||||
await self.store.save(job)
|
||||
await self._notify_progress(job)
|
||||
|
||||
return result
|
||||
|
||||
async def _notify_progress(self, job: Job):
|
||||
"""Notify progress callback if registered."""
|
||||
callback = self._progress_callbacks.get(job.id)
|
||||
if callback:
|
||||
try:
|
||||
if asyncio.iscoroutinefunction(callback):
|
||||
await callback(job)
|
||||
else:
|
||||
callback(job)
|
||||
except Exception as e:
|
||||
logger.warning(f"Progress callback error for {job.id}: {e}")
|
||||
|
||||
async def _send_webhook(self, job: Job):
|
||||
"""Send webhook notification for job completion."""
|
||||
if not job.webhook_url:
|
||||
return
|
||||
|
||||
try:
|
||||
import httpx
|
||||
|
||||
payload = {
|
||||
"job_id": job.id,
|
||||
"status": job.status.value,
|
||||
"result": job.result if job.status == JobStatus.SUCCEEDED else None,
|
||||
"error": job.error if job.status == JobStatus.FAILED else None,
|
||||
"completed_at": job.completed_at.isoformat() if job.completed_at else None,
|
||||
"duration_seconds": job.duration_seconds
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(
|
||||
job.webhook_url,
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
logger.warning(f"Webhook failed for {job.id}: {response.status_code}")
|
||||
else:
|
||||
logger.info(f"Webhook sent for {job.id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Webhook error for {job.id}: {e}")
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get executor statistics."""
|
||||
return {
|
||||
"running_jobs": len(self._running_tasks),
|
||||
"max_concurrent": self.max_concurrent,
|
||||
"default_timeout": self.default_timeout,
|
||||
"shutdown": self._shutdown
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
"""
|
||||
Job Models for PraisonAI Async Jobs API.
|
||||
|
||||
Defines the data structures for job submission, status, and results.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from typing import Optional, Dict, Any, List
|
||||
from datetime import datetime
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class JobStatus(str, Enum):
|
||||
"""Status of a job."""
|
||||
QUEUED = "queued"
|
||||
RUNNING = "running"
|
||||
SUCCEEDED = "succeeded"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
class JobSubmitRequest(BaseModel):
|
||||
"""Request body for submitting a new job."""
|
||||
prompt: str = Field(..., description="The prompt or task for the agent")
|
||||
agent_file: Optional[str] = Field(None, description="Path to agents.yaml file")
|
||||
agent_yaml: Optional[str] = Field(None, description="Inline agent YAML configuration")
|
||||
framework: Optional[str] = Field("praisonai", description="Framework to use (praisonai, crewai, autogen)")
|
||||
config: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Additional configuration")
|
||||
webhook_url: Optional[str] = Field(None, description="URL to POST results when complete")
|
||||
timeout: Optional[int] = Field(3600, description="Timeout in seconds (default: 1 hour)")
|
||||
session_id: Optional[str] = Field(None, description="Session ID for conversation continuity")
|
||||
idempotency_key: Optional[str] = Field(None, description="Idempotency key to prevent duplicate submissions")
|
||||
|
||||
|
||||
class JobSubmitResponse(BaseModel):
|
||||
"""Response after submitting a job."""
|
||||
job_id: str = Field(..., description="Unique job identifier")
|
||||
status: JobStatus = Field(..., description="Current job status")
|
||||
created_at: datetime = Field(..., description="Job creation timestamp")
|
||||
poll_url: str = Field(..., description="URL to poll for status")
|
||||
stream_url: str = Field(..., description="URL for SSE streaming")
|
||||
|
||||
|
||||
class JobProgress(BaseModel):
|
||||
"""Progress information for a running job."""
|
||||
percentage: float = Field(0.0, ge=0.0, le=100.0, description="Completion percentage")
|
||||
current_step: Optional[str] = Field(None, description="Current step description")
|
||||
steps_completed: int = Field(0, description="Number of steps completed")
|
||||
steps_total: Optional[int] = Field(None, description="Total number of steps")
|
||||
|
||||
|
||||
class JobStatusResponse(BaseModel):
|
||||
"""Response for job status query."""
|
||||
job_id: str = Field(..., description="Unique job identifier")
|
||||
status: JobStatus = Field(..., description="Current job status")
|
||||
progress: JobProgress = Field(default_factory=JobProgress, description="Progress information")
|
||||
created_at: datetime = Field(..., description="Job creation timestamp")
|
||||
started_at: Optional[datetime] = Field(None, description="When job started running")
|
||||
completed_at: Optional[datetime] = Field(None, description="When job completed")
|
||||
agent_id: Optional[str] = Field(None, description="Current agent ID")
|
||||
run_id: Optional[str] = Field(None, description="Run ID for tracing")
|
||||
session_id: Optional[str] = Field(None, description="Session ID")
|
||||
error: Optional[str] = Field(None, description="Error message if failed")
|
||||
retry_after: Optional[int] = Field(None, description="Suggested seconds before next poll")
|
||||
|
||||
|
||||
class JobResultResponse(BaseModel):
|
||||
"""Response for job result query."""
|
||||
job_id: str = Field(..., description="Unique job identifier")
|
||||
status: JobStatus = Field(..., description="Final job status")
|
||||
result: Optional[Any] = Field(None, description="Job result/output")
|
||||
result_url: Optional[str] = Field(None, description="URL to fetch large results")
|
||||
metrics: Optional[Dict[str, Any]] = Field(None, description="Execution metrics")
|
||||
created_at: datetime = Field(..., description="Job creation timestamp")
|
||||
started_at: Optional[datetime] = Field(None, description="When job started running")
|
||||
completed_at: Optional[datetime] = Field(None, description="When job completed")
|
||||
duration_seconds: Optional[float] = Field(None, description="Total execution time")
|
||||
error: Optional[str] = Field(None, description="Error message if failed")
|
||||
|
||||
|
||||
class JobListResponse(BaseModel):
|
||||
"""Response for listing jobs."""
|
||||
jobs: List[JobStatusResponse] = Field(..., description="List of jobs")
|
||||
total: int = Field(..., description="Total number of jobs matching filter")
|
||||
page: int = Field(1, description="Current page number")
|
||||
page_size: int = Field(20, description="Number of jobs per page")
|
||||
|
||||
|
||||
class Job(BaseModel):
|
||||
"""Internal job representation with full state."""
|
||||
id: str = Field(default_factory=lambda: f"run_{uuid.uuid4().hex[:12]}")
|
||||
status: JobStatus = Field(default=JobStatus.QUEUED)
|
||||
|
||||
# Request data
|
||||
prompt: str = Field(...)
|
||||
agent_file: Optional[str] = Field(None)
|
||||
agent_yaml: Optional[str] = Field(None)
|
||||
framework: str = Field("praisonai")
|
||||
config: Dict[str, Any] = Field(default_factory=dict)
|
||||
webhook_url: Optional[str] = Field(None)
|
||||
timeout: int = Field(3600)
|
||||
session_id: Optional[str] = Field(None)
|
||||
idempotency_key: Optional[str] = Field(None)
|
||||
|
||||
# Progress tracking
|
||||
progress_percentage: float = Field(0.0)
|
||||
progress_step: Optional[str] = Field(None)
|
||||
steps_completed: int = Field(0)
|
||||
steps_total: Optional[int] = Field(None)
|
||||
|
||||
# Attribution
|
||||
agent_id: Optional[str] = Field(None)
|
||||
run_id: Optional[str] = Field(None)
|
||||
|
||||
# Timestamps
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
started_at: Optional[datetime] = Field(None)
|
||||
completed_at: Optional[datetime] = Field(None)
|
||||
|
||||
# Result
|
||||
result: Optional[Any] = Field(None)
|
||||
error: Optional[str] = Field(None)
|
||||
metrics: Optional[Dict[str, Any]] = Field(None)
|
||||
|
||||
# Internal
|
||||
_cancel_requested: bool = False
|
||||
|
||||
@property
|
||||
def duration_seconds(self) -> Optional[float]:
|
||||
"""Calculate job duration."""
|
||||
if self.started_at is None:
|
||||
return None
|
||||
end_time = self.completed_at or datetime.utcnow()
|
||||
return (end_time - self.started_at).total_seconds()
|
||||
|
||||
@property
|
||||
def is_terminal(self) -> bool:
|
||||
"""Check if job is in a terminal state."""
|
||||
return self.status in (JobStatus.SUCCEEDED, JobStatus.FAILED, JobStatus.CANCELLED)
|
||||
|
||||
def to_status_response(self, base_url: str = "") -> JobStatusResponse:
|
||||
"""Convert to status response."""
|
||||
retry_after = None
|
||||
if self.status == JobStatus.QUEUED:
|
||||
retry_after = 2
|
||||
elif self.status == JobStatus.RUNNING:
|
||||
retry_after = 5
|
||||
|
||||
return JobStatusResponse(
|
||||
job_id=self.id,
|
||||
status=self.status,
|
||||
progress=JobProgress(
|
||||
percentage=self.progress_percentage,
|
||||
current_step=self.progress_step,
|
||||
steps_completed=self.steps_completed,
|
||||
steps_total=self.steps_total
|
||||
),
|
||||
created_at=self.created_at,
|
||||
started_at=self.started_at,
|
||||
completed_at=self.completed_at,
|
||||
agent_id=self.agent_id,
|
||||
run_id=self.run_id,
|
||||
session_id=self.session_id,
|
||||
error=self.error,
|
||||
retry_after=retry_after
|
||||
)
|
||||
|
||||
def to_result_response(self) -> JobResultResponse:
|
||||
"""Convert to result response."""
|
||||
return JobResultResponse(
|
||||
job_id=self.id,
|
||||
status=self.status,
|
||||
result=self.result,
|
||||
result_url=None, # For large results, would be set
|
||||
metrics=self.metrics,
|
||||
created_at=self.created_at,
|
||||
started_at=self.started_at,
|
||||
completed_at=self.completed_at,
|
||||
duration_seconds=self.duration_seconds,
|
||||
error=self.error
|
||||
)
|
||||
|
||||
def start(self):
|
||||
"""Mark job as started."""
|
||||
self.status = JobStatus.RUNNING
|
||||
self.started_at = datetime.utcnow()
|
||||
|
||||
def succeed(self, result: Any, metrics: Optional[Dict[str, Any]] = None):
|
||||
"""Mark job as succeeded."""
|
||||
self.status = JobStatus.SUCCEEDED
|
||||
self.result = result
|
||||
self.metrics = metrics
|
||||
self.completed_at = datetime.utcnow()
|
||||
self.progress_percentage = 100.0
|
||||
|
||||
def fail(self, error: str):
|
||||
"""Mark job as failed."""
|
||||
self.status = JobStatus.FAILED
|
||||
self.error = error
|
||||
self.completed_at = datetime.utcnow()
|
||||
|
||||
def cancel(self):
|
||||
"""Mark job as cancelled."""
|
||||
self.status = JobStatus.CANCELLED
|
||||
self.completed_at = datetime.utcnow()
|
||||
self._cancel_requested = True
|
||||
|
||||
def update_progress(
|
||||
self,
|
||||
percentage: Optional[float] = None,
|
||||
step: Optional[str] = None,
|
||||
steps_completed: Optional[int] = None,
|
||||
steps_total: Optional[int] = None
|
||||
):
|
||||
"""Update job progress."""
|
||||
if percentage is not None:
|
||||
self.progress_percentage = min(max(percentage, 0.0), 100.0)
|
||||
if step is not None:
|
||||
self.progress_step = step
|
||||
if steps_completed is not None:
|
||||
self.steps_completed = steps_completed
|
||||
if steps_total is not None:
|
||||
self.steps_total = steps_total
|
||||
@@ -0,0 +1,314 @@
|
||||
"""
|
||||
FastAPI Router for PraisonAI Async Jobs API.
|
||||
|
||||
Provides HTTP endpoints for job management.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Header, Request, Query, Response
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
from .models import (
|
||||
Job,
|
||||
JobStatus,
|
||||
JobSubmitRequest,
|
||||
JobSubmitResponse,
|
||||
JobStatusResponse,
|
||||
JobResultResponse,
|
||||
JobListResponse
|
||||
)
|
||||
from .store import JobStore
|
||||
from .executor import JobExecutor
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_router(store: JobStore, executor: JobExecutor) -> APIRouter:
|
||||
"""
|
||||
Create the jobs API router.
|
||||
|
||||
Args:
|
||||
store: Job storage backend
|
||||
executor: Job executor
|
||||
|
||||
Returns:
|
||||
FastAPI router with all endpoints
|
||||
"""
|
||||
router = APIRouter(prefix="/api/v1/runs", tags=["jobs"])
|
||||
|
||||
@router.post("", response_model=JobSubmitResponse, status_code=202)
|
||||
async def submit_job(
|
||||
request: Request,
|
||||
response: Response,
|
||||
body: JobSubmitRequest,
|
||||
idempotency_key: Optional[str] = Header(None, alias="Idempotency-Key")
|
||||
):
|
||||
"""
|
||||
Submit a new job for execution.
|
||||
|
||||
Returns 202 Accepted with job ID and status URLs.
|
||||
Headers returned:
|
||||
- Location: URL to poll for status
|
||||
- Retry-After: Suggested seconds before first poll (default: 2)
|
||||
|
||||
If an Idempotency-Key header is provided and a job with that key
|
||||
already exists, returns the existing job instead of creating a new one.
|
||||
"""
|
||||
# Check idempotency key
|
||||
effective_key = idempotency_key or body.idempotency_key
|
||||
if effective_key:
|
||||
existing = await store.get_by_idempotency_key(effective_key)
|
||||
if existing:
|
||||
logger.info(f"Returning existing job for idempotency key: {effective_key}")
|
||||
base_url = str(request.base_url).rstrip("/")
|
||||
poll_url = f"{base_url}/api/v1/runs/{existing.id}"
|
||||
response.headers["Location"] = poll_url
|
||||
response.headers["Retry-After"] = "2"
|
||||
return JobSubmitResponse(
|
||||
job_id=existing.id,
|
||||
status=existing.status,
|
||||
created_at=existing.created_at,
|
||||
poll_url=poll_url,
|
||||
stream_url=f"{base_url}/api/v1/runs/{existing.id}/stream"
|
||||
)
|
||||
|
||||
# Create new job
|
||||
job = Job(
|
||||
prompt=body.prompt,
|
||||
agent_file=body.agent_file,
|
||||
agent_yaml=body.agent_yaml,
|
||||
framework=body.framework or "praisonai",
|
||||
config=body.config or {},
|
||||
webhook_url=body.webhook_url,
|
||||
timeout=body.timeout or 3600,
|
||||
session_id=body.session_id,
|
||||
idempotency_key=effective_key
|
||||
)
|
||||
|
||||
# Submit for execution
|
||||
await executor.submit(job)
|
||||
|
||||
# Build response with Location and Retry-After headers (RFC best practice)
|
||||
base_url = str(request.base_url).rstrip("/")
|
||||
poll_url = f"{base_url}/api/v1/runs/{job.id}"
|
||||
response.headers["Location"] = poll_url
|
||||
response.headers["Retry-After"] = "2"
|
||||
|
||||
return JobSubmitResponse(
|
||||
job_id=job.id,
|
||||
status=job.status,
|
||||
created_at=job.created_at,
|
||||
poll_url=poll_url,
|
||||
stream_url=f"{base_url}/api/v1/runs/{job.id}/stream"
|
||||
)
|
||||
|
||||
@router.get("", response_model=JobListResponse)
|
||||
async def list_jobs(
|
||||
status: Optional[str] = Query(None, description="Filter by status"),
|
||||
session_id: Optional[str] = Query(None, description="Filter by session ID"),
|
||||
page: int = Query(1, ge=1, description="Page number"),
|
||||
page_size: int = Query(20, ge=1, le=100, description="Jobs per page")
|
||||
):
|
||||
"""
|
||||
List jobs with optional filters.
|
||||
"""
|
||||
# Parse status filter
|
||||
status_filter = None
|
||||
if status:
|
||||
try:
|
||||
status_filter = JobStatus(status)
|
||||
except ValueError:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Invalid status: {status}. Valid values: {[s.value for s in JobStatus]}"
|
||||
)
|
||||
|
||||
# Get jobs
|
||||
offset = (page - 1) * page_size
|
||||
jobs = await store.list_jobs(
|
||||
status=status_filter,
|
||||
session_id=session_id,
|
||||
limit=page_size,
|
||||
offset=offset
|
||||
)
|
||||
|
||||
total = await store.count(status=status_filter, session_id=session_id)
|
||||
|
||||
return JobListResponse(
|
||||
jobs=[job.to_status_response() for job in jobs],
|
||||
total=total,
|
||||
page=page,
|
||||
page_size=page_size
|
||||
)
|
||||
|
||||
@router.get("/{job_id}", response_model=JobStatusResponse)
|
||||
async def get_job_status(job_id: str):
|
||||
"""
|
||||
Get the status of a job.
|
||||
|
||||
Returns current status, progress, and timing information.
|
||||
"""
|
||||
job = await store.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
|
||||
|
||||
return job.to_status_response()
|
||||
|
||||
@router.get("/{job_id}/result", response_model=JobResultResponse)
|
||||
async def get_job_result(job_id: str):
|
||||
"""
|
||||
Get the result of a completed job.
|
||||
|
||||
Returns 409 Conflict if job is not yet complete.
|
||||
"""
|
||||
job = await store.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
|
||||
|
||||
if not job.is_terminal:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=f"Job is not complete. Current status: {job.status.value}"
|
||||
)
|
||||
|
||||
return job.to_result_response()
|
||||
|
||||
@router.post("/{job_id}/cancel", response_model=JobStatusResponse)
|
||||
async def cancel_job(job_id: str):
|
||||
"""
|
||||
Cancel a running job.
|
||||
|
||||
Returns 409 Conflict if job is already complete.
|
||||
"""
|
||||
job = await store.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
|
||||
|
||||
if job.is_terminal:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=f"Job is already complete. Status: {job.status.value}"
|
||||
)
|
||||
|
||||
success = await executor.cancel(job_id)
|
||||
if not success:
|
||||
raise HTTPException(status_code=500, detail="Failed to cancel job")
|
||||
|
||||
# Refresh job state
|
||||
job = await store.get(job_id)
|
||||
return job.to_status_response()
|
||||
|
||||
@router.delete("/{job_id}", status_code=204)
|
||||
async def delete_job(job_id: str):
|
||||
"""
|
||||
Delete a job.
|
||||
|
||||
Only completed jobs can be deleted.
|
||||
"""
|
||||
job = await store.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
|
||||
|
||||
if not job.is_terminal:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Cannot delete a running job. Cancel it first."
|
||||
)
|
||||
|
||||
await store.delete(job_id)
|
||||
|
||||
@router.get("/{job_id}/stream")
|
||||
async def stream_job(job_id: str):
|
||||
"""
|
||||
Stream job progress via Server-Sent Events (SSE).
|
||||
|
||||
Events:
|
||||
- status: Job status updates
|
||||
- progress: Progress percentage updates
|
||||
- result: Final result (when complete)
|
||||
- error: Error message (when failed)
|
||||
"""
|
||||
job = await store.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
|
||||
|
||||
async def event_generator():
|
||||
"""Generate SSE events for job progress."""
|
||||
last_status = None
|
||||
last_progress = -1
|
||||
|
||||
# Register for progress updates
|
||||
progress_queue = asyncio.Queue()
|
||||
|
||||
async def on_progress(updated_job: Job):
|
||||
await progress_queue.put(updated_job)
|
||||
|
||||
executor.register_progress_callback(job_id, on_progress)
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Get current job state
|
||||
current_job = await store.get(job_id)
|
||||
if not current_job:
|
||||
yield {
|
||||
"event": "error",
|
||||
"data": '{"error": "Job not found"}'
|
||||
}
|
||||
break
|
||||
|
||||
# Send status update if changed
|
||||
if current_job.status != last_status:
|
||||
last_status = current_job.status
|
||||
yield {
|
||||
"event": "status",
|
||||
"data": f'{{"status": "{current_job.status.value}", "job_id": "{job_id}"}}'
|
||||
}
|
||||
|
||||
# Send progress update if changed
|
||||
if current_job.progress_percentage != last_progress:
|
||||
last_progress = current_job.progress_percentage
|
||||
yield {
|
||||
"event": "progress",
|
||||
"data": f'{{"percentage": {current_job.progress_percentage}, "step": "{current_job.progress_step or ""}"}}'
|
||||
}
|
||||
|
||||
# Check if complete
|
||||
if current_job.is_terminal:
|
||||
if current_job.status == JobStatus.SUCCEEDED:
|
||||
result_str = str(current_job.result).replace('"', '\\"')[:1000]
|
||||
yield {
|
||||
"event": "result",
|
||||
"data": f'{{"result": "{result_str}"}}'
|
||||
}
|
||||
elif current_job.status == JobStatus.FAILED:
|
||||
error_str = (current_job.error or "Unknown error").replace('"', '\\"')
|
||||
yield {
|
||||
"event": "error",
|
||||
"data": f'{{"error": "{error_str}"}}'
|
||||
}
|
||||
elif current_job.status == JobStatus.CANCELLED:
|
||||
yield {
|
||||
"event": "cancelled",
|
||||
"data": '{"message": "Job was cancelled"}'
|
||||
}
|
||||
break
|
||||
|
||||
# Wait for next update or timeout
|
||||
try:
|
||||
await asyncio.wait_for(progress_queue.get(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
# Send heartbeat
|
||||
yield {
|
||||
"event": "heartbeat",
|
||||
"data": f'{{"timestamp": "{datetime.utcnow().isoformat()}"}}'
|
||||
}
|
||||
finally:
|
||||
executor.unregister_progress_callback(job_id)
|
||||
|
||||
return EventSourceResponse(event_generator())
|
||||
|
||||
return router
|
||||
@@ -0,0 +1,186 @@
|
||||
"""
|
||||
Server for PraisonAI Async Jobs API.
|
||||
|
||||
Provides FastAPI application setup and server startup.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from .store import InMemoryJobStore, JobStore
|
||||
from .executor import JobExecutor
|
||||
from .router import create_router
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global instances (for single-process deployment)
|
||||
_store: Optional[JobStore] = None
|
||||
_executor: Optional[JobExecutor] = None
|
||||
|
||||
|
||||
def get_store() -> JobStore:
|
||||
"""Get or create the job store."""
|
||||
global _store
|
||||
if _store is None:
|
||||
_store = InMemoryJobStore(max_jobs=1000)
|
||||
return _store
|
||||
|
||||
|
||||
def get_executor() -> JobExecutor:
|
||||
"""Get or create the job executor."""
|
||||
global _executor
|
||||
if _executor is None:
|
||||
_executor = JobExecutor(
|
||||
store=get_store(),
|
||||
max_concurrent=int(os.environ.get("PRAISONAI_MAX_CONCURRENT_JOBS", "10")),
|
||||
default_timeout=int(os.environ.get("PRAISONAI_JOB_TIMEOUT", "3600"))
|
||||
)
|
||||
return _executor
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Application lifespan manager."""
|
||||
executor = get_executor()
|
||||
await executor.start()
|
||||
logger.info("Jobs API server started")
|
||||
|
||||
yield
|
||||
|
||||
await executor.stop()
|
||||
logger.info("Jobs API server stopped")
|
||||
|
||||
|
||||
def create_app(
|
||||
store: Optional[JobStore] = None,
|
||||
executor: Optional[JobExecutor] = None,
|
||||
cors_origins: Optional[list] = None
|
||||
) -> FastAPI:
|
||||
"""
|
||||
Create the FastAPI application.
|
||||
|
||||
Args:
|
||||
store: Optional custom job store
|
||||
executor: Optional custom executor
|
||||
cors_origins: Optional list of allowed CORS origins
|
||||
|
||||
Returns:
|
||||
FastAPI application
|
||||
"""
|
||||
global _store, _executor
|
||||
|
||||
if store:
|
||||
_store = store
|
||||
if executor:
|
||||
_executor = executor
|
||||
|
||||
app = FastAPI(
|
||||
title="PraisonAI Jobs API",
|
||||
description="Async Jobs API for long-running agent tasks",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
# Add CORS middleware
|
||||
origins = cors_origins or ["*"]
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=origins,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Add jobs router
|
||||
jobs_router = create_router(get_store(), get_executor())
|
||||
app.include_router(jobs_router)
|
||||
|
||||
# Health check endpoint
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
"""Health check endpoint."""
|
||||
return {
|
||||
"status": "healthy",
|
||||
"store": get_store().__class__.__name__,
|
||||
"executor_stats": get_executor().get_stats()
|
||||
}
|
||||
|
||||
# Stats endpoint
|
||||
@app.get("/stats")
|
||||
async def get_stats():
|
||||
"""Get server statistics."""
|
||||
store = get_store()
|
||||
executor = get_executor()
|
||||
|
||||
stats = {
|
||||
"executor": executor.get_stats()
|
||||
}
|
||||
|
||||
if hasattr(store, "get_stats"):
|
||||
stats["store"] = store.get_stats()
|
||||
|
||||
return stats
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def start_server(
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8005,
|
||||
reload: bool = False,
|
||||
workers: int = 1,
|
||||
log_level: str = "info"
|
||||
):
|
||||
"""
|
||||
Start the jobs API server.
|
||||
|
||||
Args:
|
||||
host: Host to bind to
|
||||
port: Port to bind to
|
||||
reload: Enable auto-reload for development
|
||||
workers: Number of worker processes
|
||||
log_level: Logging level
|
||||
"""
|
||||
try:
|
||||
import uvicorn
|
||||
except ImportError:
|
||||
raise RuntimeError("uvicorn is required. Install with: pip install uvicorn")
|
||||
|
||||
logger.info(f"Starting PraisonAI Jobs API on {host}:{port}")
|
||||
|
||||
uvicorn.run(
|
||||
"praisonai.jobs.server:create_app",
|
||||
host=host,
|
||||
port=port,
|
||||
reload=reload,
|
||||
workers=workers,
|
||||
log_level=log_level,
|
||||
factory=True
|
||||
)
|
||||
|
||||
|
||||
# For direct module execution
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="PraisonAI Jobs API Server")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Host to bind to")
|
||||
parser.add_argument("--port", type=int, default=8005, help="Port to bind to")
|
||||
parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
|
||||
parser.add_argument("--workers", type=int, default=1, help="Number of workers")
|
||||
parser.add_argument("--log-level", default="info", help="Log level")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
start_server(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
reload=args.reload,
|
||||
workers=args.workers,
|
||||
log_level=args.log_level
|
||||
)
|
||||
@@ -0,0 +1,203 @@
|
||||
"""
|
||||
Job Store for PraisonAI Async Jobs API.
|
||||
|
||||
Provides storage backends for job state persistence.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from .models import Job, JobStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class JobStore(ABC):
|
||||
"""Abstract base class for job storage backends."""
|
||||
|
||||
@abstractmethod
|
||||
async def save(self, job: Job) -> None:
|
||||
"""Save or update a job."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get(self, job_id: str) -> Optional[Job]:
|
||||
"""Get a job by ID."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_by_idempotency_key(self, key: str) -> Optional[Job]:
|
||||
"""Get a job by idempotency key."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def list_jobs(
|
||||
self,
|
||||
status: Optional[JobStatus] = None,
|
||||
session_id: Optional[str] = None,
|
||||
limit: int = 20,
|
||||
offset: int = 0
|
||||
) -> List[Job]:
|
||||
"""List jobs with optional filters."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def count(
|
||||
self,
|
||||
status: Optional[JobStatus] = None,
|
||||
session_id: Optional[str] = None
|
||||
) -> int:
|
||||
"""Count jobs matching filters."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, job_id: str) -> bool:
|
||||
"""Delete a job by ID."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def cleanup_old_jobs(self, max_age_seconds: int = 86400) -> int:
|
||||
"""Remove jobs older than max_age_seconds. Returns count deleted."""
|
||||
pass
|
||||
|
||||
|
||||
class InMemoryJobStore(JobStore):
|
||||
"""
|
||||
In-memory job store.
|
||||
|
||||
Suitable for development and single-process deployments.
|
||||
Jobs are lost on restart.
|
||||
"""
|
||||
|
||||
def __init__(self, max_jobs: int = 1000):
|
||||
self._jobs: Dict[str, Job] = {}
|
||||
self._idempotency_keys: Dict[str, str] = {} # key -> job_id
|
||||
self._max_jobs = max_jobs
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def save(self, job: Job) -> None:
|
||||
"""Save or update a job."""
|
||||
async with self._lock:
|
||||
self._jobs[job.id] = job
|
||||
|
||||
# Track idempotency key
|
||||
if job.idempotency_key:
|
||||
self._idempotency_keys[job.idempotency_key] = job.id
|
||||
|
||||
# Enforce max jobs limit by removing oldest completed
|
||||
if len(self._jobs) > self._max_jobs:
|
||||
await self._evict_oldest_completed()
|
||||
|
||||
async def _evict_oldest_completed(self) -> None:
|
||||
"""Remove oldest completed jobs to stay under limit."""
|
||||
completed = [
|
||||
(job_id, job) for job_id, job in self._jobs.items()
|
||||
if job.is_terminal
|
||||
]
|
||||
|
||||
if not completed:
|
||||
return
|
||||
|
||||
# Sort by completed_at, oldest first
|
||||
completed.sort(key=lambda x: x[1].completed_at or datetime.min)
|
||||
|
||||
# Remove oldest 10%
|
||||
to_remove = max(1, len(completed) // 10)
|
||||
for job_id, job in completed[:to_remove]:
|
||||
del self._jobs[job_id]
|
||||
if job.idempotency_key:
|
||||
self._idempotency_keys.pop(job.idempotency_key, None)
|
||||
|
||||
async def get(self, job_id: str) -> Optional[Job]:
|
||||
"""Get a job by ID."""
|
||||
return self._jobs.get(job_id)
|
||||
|
||||
async def get_by_idempotency_key(self, key: str) -> Optional[Job]:
|
||||
"""Get a job by idempotency key."""
|
||||
job_id = self._idempotency_keys.get(key)
|
||||
if job_id:
|
||||
return self._jobs.get(job_id)
|
||||
return None
|
||||
|
||||
async def list_jobs(
|
||||
self,
|
||||
status: Optional[JobStatus] = None,
|
||||
session_id: Optional[str] = None,
|
||||
limit: int = 20,
|
||||
offset: int = 0
|
||||
) -> List[Job]:
|
||||
"""List jobs with optional filters."""
|
||||
jobs = list(self._jobs.values())
|
||||
|
||||
# Apply filters
|
||||
if status:
|
||||
jobs = [j for j in jobs if j.status == status]
|
||||
if session_id:
|
||||
jobs = [j for j in jobs if j.session_id == session_id]
|
||||
|
||||
# Sort by created_at descending (newest first)
|
||||
jobs.sort(key=lambda j: j.created_at, reverse=True)
|
||||
|
||||
# Apply pagination
|
||||
return jobs[offset:offset + limit]
|
||||
|
||||
async def count(
|
||||
self,
|
||||
status: Optional[JobStatus] = None,
|
||||
session_id: Optional[str] = None
|
||||
) -> int:
|
||||
"""Count jobs matching filters."""
|
||||
jobs = list(self._jobs.values())
|
||||
|
||||
if status:
|
||||
jobs = [j for j in jobs if j.status == status]
|
||||
if session_id:
|
||||
jobs = [j for j in jobs if j.session_id == session_id]
|
||||
|
||||
return len(jobs)
|
||||
|
||||
async def delete(self, job_id: str) -> bool:
|
||||
"""Delete a job by ID."""
|
||||
async with self._lock:
|
||||
job = self._jobs.pop(job_id, None)
|
||||
if job:
|
||||
if job.idempotency_key:
|
||||
self._idempotency_keys.pop(job.idempotency_key, None)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def cleanup_old_jobs(self, max_age_seconds: int = 86400) -> int:
|
||||
"""Remove jobs older than max_age_seconds."""
|
||||
async with self._lock:
|
||||
cutoff = datetime.utcnow() - timedelta(seconds=max_age_seconds)
|
||||
to_remove = []
|
||||
|
||||
for job_id, job in self._jobs.items():
|
||||
if job.is_terminal and job.completed_at and job.completed_at < cutoff:
|
||||
to_remove.append(job_id)
|
||||
|
||||
for job_id in to_remove:
|
||||
job = self._jobs.pop(job_id)
|
||||
if job.idempotency_key:
|
||||
self._idempotency_keys.pop(job.idempotency_key, None)
|
||||
|
||||
if to_remove:
|
||||
logger.info(f"Cleaned up {len(to_remove)} old jobs")
|
||||
|
||||
return len(to_remove)
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get store statistics."""
|
||||
status_counts = {}
|
||||
for job in self._jobs.values():
|
||||
status_counts[job.status.value] = status_counts.get(job.status.value, 0) + 1
|
||||
|
||||
return {
|
||||
"total_jobs": len(self._jobs),
|
||||
"idempotency_keys": len(self._idempotency_keys),
|
||||
"max_jobs": self._max_jobs,
|
||||
"status_counts": status_counts
|
||||
}
|
||||
@@ -1 +1 @@
|
||||
__version__ = "2.5.0"
|
||||
__version__ = "2.5.1"
|
||||
@@ -12,7 +12,7 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=0.3.0",
|
||||
"praisonaiagents>=0.3.1",
|
||||
"python-dotenv>=0.19.0",
|
||||
"instructor>=1.3.3",
|
||||
"PyYAML>=6.0",
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Integration tests for knowledge stack."""
|
||||
@@ -0,0 +1,308 @@
|
||||
"""
|
||||
Integration tests for Knowledge Stack.
|
||||
|
||||
These tests require external services and are skipped by default.
|
||||
Set environment variables to enable:
|
||||
- OPENAI_API_KEY: For embedding tests
|
||||
- PINECONE_API_KEY: For Pinecone tests
|
||||
- COHERE_API_KEY: For Cohere reranker tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
import tempfile
|
||||
|
||||
# Skip markers
|
||||
requires_openai = pytest.mark.skipif(
|
||||
not os.environ.get("OPENAI_API_KEY"),
|
||||
reason="OPENAI_API_KEY not set"
|
||||
)
|
||||
|
||||
requires_pinecone = pytest.mark.skipif(
|
||||
not os.environ.get("PINECONE_API_KEY"),
|
||||
reason="PINECONE_API_KEY not set"
|
||||
)
|
||||
|
||||
requires_cohere = pytest.mark.skipif(
|
||||
not os.environ.get("COHERE_API_KEY"),
|
||||
reason="COHERE_API_KEY not set"
|
||||
)
|
||||
|
||||
requires_chromadb = pytest.mark.skipif(
|
||||
not pytest.importorskip("chromadb", reason="chromadb not installed"),
|
||||
reason="chromadb not installed"
|
||||
)
|
||||
|
||||
|
||||
class TestKnowledgeIntegration:
|
||||
"""Integration tests for Knowledge class."""
|
||||
|
||||
@requires_openai
|
||||
def test_knowledge_add_and_search(self):
|
||||
"""Test adding and searching documents."""
|
||||
from praisonaiagents import Knowledge
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Create a test file
|
||||
test_file = os.path.join(tmpdir, "test.txt")
|
||||
with open(test_file, "w") as f:
|
||||
f.write("Python is a programming language created by Guido van Rossum.")
|
||||
|
||||
# Create knowledge instance
|
||||
knowledge = Knowledge(config={
|
||||
"vector_store": {
|
||||
"provider": "chroma",
|
||||
"config": {
|
||||
"path": os.path.join(tmpdir, ".praison"),
|
||||
"collection_name": "test_collection"
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
# Add document
|
||||
result = knowledge.add(test_file)
|
||||
assert result is not None
|
||||
|
||||
# Search
|
||||
results = knowledge.search("Who created Python?")
|
||||
assert len(results) > 0
|
||||
|
||||
@requires_openai
|
||||
def test_knowledge_store_and_retrieve(self):
|
||||
"""Test storing and retrieving text content."""
|
||||
from praisonaiagents import Knowledge
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
knowledge = Knowledge(config={
|
||||
"vector_store": {
|
||||
"provider": "chroma",
|
||||
"config": {
|
||||
"path": os.path.join(tmpdir, ".praison"),
|
||||
"collection_name": "test_store"
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
# Store text
|
||||
knowledge.store("The capital of France is Paris.")
|
||||
knowledge.store("The capital of Germany is Berlin.")
|
||||
|
||||
# Search
|
||||
results = knowledge.search("What is the capital of France?")
|
||||
assert len(results) > 0
|
||||
# Check that Paris is mentioned in results
|
||||
result_text = str(results[0]).lower()
|
||||
assert "paris" in result_text or "france" in result_text
|
||||
|
||||
|
||||
class TestVectorStoreIntegration:
|
||||
"""Integration tests for vector stores."""
|
||||
|
||||
@requires_openai
|
||||
def test_chroma_vector_store(self):
|
||||
"""Test ChromaDB vector store."""
|
||||
try:
|
||||
from praisonai.adapters.vector_stores import ChromaVectorStore
|
||||
except ImportError:
|
||||
pytest.skip("ChromaDB not available")
|
||||
|
||||
import openai
|
||||
client = openai.OpenAI()
|
||||
|
||||
def get_embedding(text):
|
||||
response = client.embeddings.create(
|
||||
input=text,
|
||||
model="text-embedding-3-small"
|
||||
)
|
||||
return response.data[0].embedding
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
store = ChromaVectorStore(
|
||||
persist_directory=tmpdir,
|
||||
namespace="test"
|
||||
)
|
||||
|
||||
# Add vectors
|
||||
texts = ["Hello world", "Goodbye world"]
|
||||
embeddings = [get_embedding(t) for t in texts]
|
||||
|
||||
ids = store.add(texts=texts, embeddings=embeddings)
|
||||
assert len(ids) == 2
|
||||
|
||||
# Query
|
||||
query_embedding = get_embedding("Hello")
|
||||
results = store.query(embedding=query_embedding, top_k=2)
|
||||
|
||||
assert len(results) > 0
|
||||
assert results[0].text == "Hello world"
|
||||
|
||||
@requires_pinecone
|
||||
def test_pinecone_vector_store(self):
|
||||
"""Test Pinecone vector store."""
|
||||
try:
|
||||
from praisonai.adapters.vector_stores import PineconeVectorStore
|
||||
except ImportError:
|
||||
pytest.skip("Pinecone not available")
|
||||
|
||||
import openai
|
||||
client = openai.OpenAI()
|
||||
|
||||
def get_embedding(text):
|
||||
response = client.embeddings.create(
|
||||
input=text,
|
||||
model="text-embedding-3-small"
|
||||
)
|
||||
return response.data[0].embedding
|
||||
|
||||
store = PineconeVectorStore(
|
||||
index_name="praisonai-test",
|
||||
namespace="integration-test"
|
||||
)
|
||||
|
||||
# Add vectors
|
||||
texts = ["Integration test document"]
|
||||
embeddings = [get_embedding(t) for t in texts]
|
||||
|
||||
ids = store.add(texts=texts, embeddings=embeddings)
|
||||
assert len(ids) == 1
|
||||
|
||||
# Query
|
||||
query_embedding = get_embedding("test document")
|
||||
results = store.query(embedding=query_embedding, top_k=1)
|
||||
|
||||
assert len(results) > 0
|
||||
|
||||
# Cleanup
|
||||
store.delete(ids=ids)
|
||||
|
||||
|
||||
class TestRerankerIntegration:
|
||||
"""Integration tests for rerankers."""
|
||||
|
||||
@requires_openai
|
||||
def test_llm_reranker(self):
|
||||
"""Test LLM-based reranker."""
|
||||
try:
|
||||
from praisonai.adapters.rerankers import LLMReranker
|
||||
except ImportError:
|
||||
pytest.skip("Reranker not available")
|
||||
|
||||
reranker = LLMReranker(model="gpt-4o-mini")
|
||||
|
||||
documents = [
|
||||
"Python is a programming language.",
|
||||
"JavaScript is used for web development.",
|
||||
"Python was created by Guido van Rossum."
|
||||
]
|
||||
|
||||
results = reranker.rerank("Who created Python?", documents, top_k=2)
|
||||
|
||||
assert len(results) == 2
|
||||
# The document about Guido should rank higher
|
||||
assert "Guido" in results[0].text
|
||||
|
||||
@requires_cohere
|
||||
def test_cohere_reranker(self):
|
||||
"""Test Cohere reranker."""
|
||||
try:
|
||||
from praisonai.adapters.rerankers import CohereReranker
|
||||
except ImportError:
|
||||
pytest.skip("Cohere not available")
|
||||
|
||||
reranker = CohereReranker()
|
||||
|
||||
documents = [
|
||||
"Python is a programming language.",
|
||||
"JavaScript is used for web development.",
|
||||
"Python was created by Guido van Rossum."
|
||||
]
|
||||
|
||||
results = reranker.rerank("Who created Python?", documents, top_k=2)
|
||||
|
||||
assert len(results) == 2
|
||||
|
||||
|
||||
class TestRetrieverIntegration:
|
||||
"""Integration tests for retrievers."""
|
||||
|
||||
@requires_openai
|
||||
def test_basic_retriever(self):
|
||||
"""Test basic retriever with real embeddings."""
|
||||
from praisonaiagents.knowledge.vector_store import InMemoryVectorStore
|
||||
try:
|
||||
from praisonai.adapters.retrievers import BasicRetriever
|
||||
except ImportError:
|
||||
pytest.skip("Retriever not available")
|
||||
|
||||
import openai
|
||||
client = openai.OpenAI()
|
||||
|
||||
def get_embedding(text):
|
||||
response = client.embeddings.create(
|
||||
input=text,
|
||||
model="text-embedding-3-small"
|
||||
)
|
||||
return response.data[0].embedding
|
||||
|
||||
# Setup store with documents
|
||||
store = InMemoryVectorStore()
|
||||
texts = [
|
||||
"Python is a programming language.",
|
||||
"JavaScript is for web development.",
|
||||
"Machine learning uses neural networks."
|
||||
]
|
||||
embeddings = [get_embedding(t) for t in texts]
|
||||
store.add(texts=texts, embeddings=embeddings)
|
||||
|
||||
# Create retriever
|
||||
retriever = BasicRetriever(
|
||||
vector_store=store,
|
||||
embedding_fn=get_embedding
|
||||
)
|
||||
|
||||
# Retrieve
|
||||
results = retriever.retrieve("What is Python?", top_k=2)
|
||||
|
||||
assert len(results) == 2
|
||||
assert "Python" in results[0].text
|
||||
|
||||
@requires_openai
|
||||
def test_fusion_retriever(self):
|
||||
"""Test fusion retriever with real embeddings."""
|
||||
from praisonaiagents.knowledge.vector_store import InMemoryVectorStore
|
||||
try:
|
||||
from praisonai.adapters.retrievers import FusionRetriever
|
||||
except ImportError:
|
||||
pytest.skip("Retriever not available")
|
||||
|
||||
import openai
|
||||
client = openai.OpenAI()
|
||||
|
||||
def get_embedding(text):
|
||||
response = client.embeddings.create(
|
||||
input=text,
|
||||
model="text-embedding-3-small"
|
||||
)
|
||||
return response.data[0].embedding
|
||||
|
||||
# Setup store with documents
|
||||
store = InMemoryVectorStore()
|
||||
texts = [
|
||||
"Python is a programming language created by Guido.",
|
||||
"JavaScript is for web development.",
|
||||
"Python is great for data science."
|
||||
]
|
||||
embeddings = [get_embedding(t) for t in texts]
|
||||
store.add(texts=texts, embeddings=embeddings)
|
||||
|
||||
# Create retriever
|
||||
retriever = FusionRetriever(
|
||||
vector_store=store,
|
||||
embedding_fn=get_embedding,
|
||||
num_queries=2
|
||||
)
|
||||
|
||||
# Retrieve
|
||||
results = retriever.retrieve("Python programming", top_k=2)
|
||||
|
||||
assert len(results) >= 1
|
||||
@@ -0,0 +1,268 @@
|
||||
"""
|
||||
Unit tests for Knowledge CLI Handler.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from praisonai.cli.features.knowledge import KnowledgeHandler
|
||||
|
||||
|
||||
class TestKnowledgeHandler:
|
||||
"""Tests for KnowledgeHandler."""
|
||||
|
||||
def test_init_defaults(self):
|
||||
"""Test default initialization."""
|
||||
handler = KnowledgeHandler()
|
||||
assert handler.vector_store == "chroma"
|
||||
assert handler.retrieval_strategy == "basic"
|
||||
assert handler.reranker is None
|
||||
assert handler.index_type == "vector"
|
||||
assert handler.query_mode == "default"
|
||||
|
||||
def test_init_custom_options(self):
|
||||
"""Test initialization with custom options."""
|
||||
handler = KnowledgeHandler(
|
||||
vector_store="pinecone",
|
||||
retrieval_strategy="fusion",
|
||||
reranker="llm",
|
||||
index_type="hybrid",
|
||||
query_mode="sub_question"
|
||||
)
|
||||
assert handler.vector_store == "pinecone"
|
||||
assert handler.retrieval_strategy == "fusion"
|
||||
assert handler.reranker == "llm"
|
||||
assert handler.index_type == "hybrid"
|
||||
assert handler.query_mode == "sub_question"
|
||||
|
||||
def test_feature_name(self):
|
||||
"""Test feature name property."""
|
||||
handler = KnowledgeHandler()
|
||||
assert handler.feature_name == "knowledge"
|
||||
|
||||
def test_get_actions(self):
|
||||
"""Test available actions."""
|
||||
handler = KnowledgeHandler()
|
||||
actions = handler.get_actions()
|
||||
assert "add" in actions
|
||||
assert "query" in actions
|
||||
assert "search" in actions
|
||||
assert "list" in actions
|
||||
assert "clear" in actions
|
||||
assert "stats" in actions
|
||||
|
||||
def test_get_help_text(self):
|
||||
"""Test help text."""
|
||||
handler = KnowledgeHandler()
|
||||
help_text = handler.get_help_text()
|
||||
assert "knowledge add" in help_text
|
||||
assert "knowledge query" in help_text
|
||||
assert "--vector-store" in help_text
|
||||
assert "--retrieval" in help_text
|
||||
|
||||
def test_action_add_no_args(self):
|
||||
"""Test add action with no arguments."""
|
||||
handler = KnowledgeHandler()
|
||||
result = handler.action_add([])
|
||||
assert result is False
|
||||
|
||||
def test_action_add_nonexistent_path(self):
|
||||
"""Test add action with non-existent path."""
|
||||
handler = KnowledgeHandler()
|
||||
with patch.object(handler, '_get_knowledge', return_value=MagicMock()):
|
||||
result = handler.action_add(["/nonexistent/path/file.txt"])
|
||||
assert result is False
|
||||
|
||||
def test_action_add_file(self):
|
||||
"""Test add action with file."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
|
||||
f.write(b"test content")
|
||||
path = f.name
|
||||
|
||||
try:
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
result = handler.action_add([path])
|
||||
assert result is True
|
||||
mock_knowledge.add.assert_called_once_with(path)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_action_add_directory(self):
|
||||
"""Test add action with directory."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Create some test files
|
||||
for i in range(3):
|
||||
with open(os.path.join(tmpdir, f"file{i}.txt"), "w") as f:
|
||||
f.write(f"content {i}")
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
result = handler.action_add([tmpdir])
|
||||
assert result is True
|
||||
assert mock_knowledge.add.call_count == 3
|
||||
|
||||
def test_action_add_url(self):
|
||||
"""Test add action with URL."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
result = handler.action_add(["https://example.com/doc.pdf"])
|
||||
assert result is True
|
||||
mock_knowledge.add.assert_called_once_with("https://example.com/doc.pdf")
|
||||
|
||||
def test_action_query_no_args(self):
|
||||
"""Test query action with no arguments."""
|
||||
handler = KnowledgeHandler()
|
||||
result = handler.action_query([])
|
||||
assert result == []
|
||||
|
||||
def test_action_query_with_results(self):
|
||||
"""Test query action with results."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
mock_knowledge.search.return_value = [
|
||||
{"memory": "Result 1", "score": 0.9},
|
||||
{"memory": "Result 2", "score": 0.8}
|
||||
]
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
results = handler.action_query(["test query"])
|
||||
assert len(results) == 2
|
||||
mock_knowledge.search.assert_called_once()
|
||||
|
||||
def test_action_search_is_alias(self):
|
||||
"""Test that search is alias for query."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
mock_knowledge.search.return_value = []
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
handler.action_search(["test"])
|
||||
mock_knowledge.search.assert_called_once()
|
||||
|
||||
def test_action_list(self):
|
||||
"""Test list action."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
mock_knowledge.documents = {"doc1": {}, "doc2": {}}
|
||||
mock_knowledge.list_documents = MagicMock(return_value=["doc1", "doc2"])
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
docs = handler.action_list([])
|
||||
assert len(docs) == 2
|
||||
|
||||
def test_action_clear(self):
|
||||
"""Test clear action."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
result = handler.action_clear([])
|
||||
assert result is True
|
||||
mock_knowledge.clear.assert_called_once()
|
||||
|
||||
def test_action_stats(self):
|
||||
"""Test stats action."""
|
||||
handler = KnowledgeHandler(
|
||||
vector_store="chroma",
|
||||
retrieval_strategy="fusion"
|
||||
)
|
||||
mock_knowledge = MagicMock()
|
||||
mock_knowledge.memory = MagicMock()
|
||||
mock_knowledge.memory.get_all.return_value = [1, 2, 3]
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
stats = handler.action_stats([])
|
||||
assert stats["vector_store"] == "chroma"
|
||||
assert stats["retrieval_strategy"] == "fusion"
|
||||
assert stats["document_count"] == 3
|
||||
|
||||
def test_action_info_is_alias(self):
|
||||
"""Test that info is alias for stats."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
mock_knowledge.memory = MagicMock()
|
||||
mock_knowledge.memory.get_all.return_value = []
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
stats = handler.action_info([])
|
||||
assert "vector_store" in stats
|
||||
|
||||
|
||||
class TestKnowledgeHandlerGlobPatterns:
|
||||
"""Tests for glob pattern handling."""
|
||||
|
||||
def test_glob_pattern_detection(self):
|
||||
"""Test that glob patterns are detected."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Create test files
|
||||
for ext in ["txt", "md", "pdf"]:
|
||||
with open(os.path.join(tmpdir, f"file.{ext}"), "w") as f:
|
||||
f.write("content")
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
pattern = os.path.join(tmpdir, "*.txt")
|
||||
result = handler.action_add([pattern])
|
||||
assert result is True
|
||||
assert mock_knowledge.add.call_count == 1
|
||||
|
||||
def test_glob_no_matches(self):
|
||||
"""Test glob pattern with no matches."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
pattern = os.path.join(tmpdir, "*.nonexistent")
|
||||
result = handler.action_add([pattern])
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestKnowledgeHandlerEdgeCases:
|
||||
"""Tests for edge cases."""
|
||||
|
||||
def test_knowledge_import_error(self):
|
||||
"""Test handling of import error."""
|
||||
handler = KnowledgeHandler()
|
||||
|
||||
with patch('praisonai.cli.features.knowledge.KnowledgeHandler._get_knowledge') as mock:
|
||||
mock.return_value = None
|
||||
result = handler.action_add(["test.txt"])
|
||||
assert result is False
|
||||
|
||||
def test_add_with_exception(self):
|
||||
"""Test add action when exception occurs."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
mock_knowledge.add.side_effect = Exception("Test error")
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
|
||||
f.write(b"test")
|
||||
path = f.name
|
||||
|
||||
try:
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
result = handler.action_add([path])
|
||||
assert result is False
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_query_with_exception(self):
|
||||
"""Test query action when exception occurs."""
|
||||
handler = KnowledgeHandler()
|
||||
mock_knowledge = MagicMock()
|
||||
mock_knowledge.search.side_effect = Exception("Test error")
|
||||
|
||||
with patch.object(handler, '_get_knowledge', return_value=mock_knowledge):
|
||||
results = handler.action_query(["test"])
|
||||
assert results == []
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for the jobs module."""
|
||||
@@ -0,0 +1,278 @@
|
||||
"""Unit tests for jobs models."""
|
||||
|
||||
from praisonai.jobs.models import (
|
||||
Job,
|
||||
JobStatus,
|
||||
JobSubmitRequest,
|
||||
JobStatusResponse,
|
||||
JobResultResponse,
|
||||
JobProgress,
|
||||
JobListResponse
|
||||
)
|
||||
|
||||
|
||||
class TestJobStatus:
|
||||
"""Tests for JobStatus enum."""
|
||||
|
||||
def test_status_values(self):
|
||||
"""Test all status values exist."""
|
||||
assert JobStatus.QUEUED.value == "queued"
|
||||
assert JobStatus.RUNNING.value == "running"
|
||||
assert JobStatus.SUCCEEDED.value == "succeeded"
|
||||
assert JobStatus.FAILED.value == "failed"
|
||||
assert JobStatus.CANCELLED.value == "cancelled"
|
||||
|
||||
def test_status_is_string_enum(self):
|
||||
"""Test status can be used as string."""
|
||||
assert str(JobStatus.RUNNING) == "JobStatus.RUNNING"
|
||||
assert JobStatus.RUNNING.value == "running"
|
||||
|
||||
|
||||
class TestJob:
|
||||
"""Tests for Job model."""
|
||||
|
||||
def test_job_creation_defaults(self):
|
||||
"""Test job creation with defaults."""
|
||||
job = Job(prompt="Test prompt")
|
||||
|
||||
assert job.prompt == "Test prompt"
|
||||
assert job.status == JobStatus.QUEUED
|
||||
assert job.id.startswith("run_")
|
||||
assert job.framework == "praisonai"
|
||||
assert job.timeout == 3600
|
||||
assert job.progress_percentage == 0.0
|
||||
assert job.created_at is not None
|
||||
assert job.started_at is None
|
||||
assert job.completed_at is None
|
||||
|
||||
def test_job_creation_with_all_fields(self):
|
||||
"""Test job creation with all fields."""
|
||||
job = Job(
|
||||
prompt="Test prompt",
|
||||
agent_file="agents.yaml",
|
||||
agent_yaml="agents:\n - name: test",
|
||||
framework="crewai",
|
||||
config={"key": "value"},
|
||||
webhook_url="https://example.com/webhook",
|
||||
timeout=7200,
|
||||
session_id="session_123",
|
||||
idempotency_key="idem_456"
|
||||
)
|
||||
|
||||
assert job.agent_file == "agents.yaml"
|
||||
assert job.agent_yaml == "agents:\n - name: test"
|
||||
assert job.framework == "crewai"
|
||||
assert job.config == {"key": "value"}
|
||||
assert job.webhook_url == "https://example.com/webhook"
|
||||
assert job.timeout == 7200
|
||||
assert job.session_id == "session_123"
|
||||
assert job.idempotency_key == "idem_456"
|
||||
|
||||
def test_job_start(self):
|
||||
"""Test job start method."""
|
||||
job = Job(prompt="Test")
|
||||
assert job.status == JobStatus.QUEUED
|
||||
assert job.started_at is None
|
||||
|
||||
job.start()
|
||||
|
||||
assert job.status == JobStatus.RUNNING
|
||||
assert job.started_at is not None
|
||||
|
||||
def test_job_succeed(self):
|
||||
"""Test job succeed method."""
|
||||
job = Job(prompt="Test")
|
||||
job.start()
|
||||
|
||||
job.succeed("Result data", {"tokens": 100})
|
||||
|
||||
assert job.status == JobStatus.SUCCEEDED
|
||||
assert job.result == "Result data"
|
||||
assert job.metrics == {"tokens": 100}
|
||||
assert job.completed_at is not None
|
||||
assert job.progress_percentage == 100.0
|
||||
|
||||
def test_job_fail(self):
|
||||
"""Test job fail method."""
|
||||
job = Job(prompt="Test")
|
||||
job.start()
|
||||
|
||||
job.fail("Something went wrong")
|
||||
|
||||
assert job.status == JobStatus.FAILED
|
||||
assert job.error == "Something went wrong"
|
||||
assert job.completed_at is not None
|
||||
|
||||
def test_job_cancel(self):
|
||||
"""Test job cancel method."""
|
||||
job = Job(prompt="Test")
|
||||
job.start()
|
||||
|
||||
job.cancel()
|
||||
|
||||
assert job.status == JobStatus.CANCELLED
|
||||
assert job.completed_at is not None
|
||||
|
||||
def test_job_is_terminal(self):
|
||||
"""Test is_terminal property."""
|
||||
job = Job(prompt="Test")
|
||||
assert not job.is_terminal
|
||||
|
||||
job.start()
|
||||
assert not job.is_terminal
|
||||
|
||||
job.succeed("Done")
|
||||
assert job.is_terminal
|
||||
|
||||
def test_job_is_terminal_failed(self):
|
||||
"""Test is_terminal for failed job."""
|
||||
job = Job(prompt="Test")
|
||||
job.start()
|
||||
job.fail("Error")
|
||||
assert job.is_terminal
|
||||
|
||||
def test_job_is_terminal_cancelled(self):
|
||||
"""Test is_terminal for cancelled job."""
|
||||
job = Job(prompt="Test")
|
||||
job.start()
|
||||
job.cancel()
|
||||
assert job.is_terminal
|
||||
|
||||
def test_job_duration_seconds(self):
|
||||
"""Test duration_seconds property."""
|
||||
job = Job(prompt="Test")
|
||||
assert job.duration_seconds is None
|
||||
|
||||
job.start()
|
||||
# Duration should be very small but not None
|
||||
assert job.duration_seconds is not None
|
||||
assert job.duration_seconds >= 0
|
||||
|
||||
def test_job_update_progress(self):
|
||||
"""Test update_progress method."""
|
||||
job = Job(prompt="Test")
|
||||
|
||||
job.update_progress(percentage=50.0, step="Processing")
|
||||
|
||||
assert job.progress_percentage == 50.0
|
||||
assert job.progress_step == "Processing"
|
||||
|
||||
def test_job_update_progress_clamps_values(self):
|
||||
"""Test update_progress clamps percentage to 0-100."""
|
||||
job = Job(prompt="Test")
|
||||
|
||||
job.update_progress(percentage=150.0)
|
||||
assert job.progress_percentage == 100.0
|
||||
|
||||
job.update_progress(percentage=-10.0)
|
||||
assert job.progress_percentage == 0.0
|
||||
|
||||
def test_job_to_status_response(self):
|
||||
"""Test to_status_response method."""
|
||||
job = Job(prompt="Test", session_id="sess_123")
|
||||
job.start()
|
||||
job.update_progress(percentage=50.0, step="Working")
|
||||
|
||||
response = job.to_status_response()
|
||||
|
||||
assert isinstance(response, JobStatusResponse)
|
||||
assert response.job_id == job.id
|
||||
assert response.status == JobStatus.RUNNING
|
||||
assert response.progress.percentage == 50.0
|
||||
assert response.progress.current_step == "Working"
|
||||
assert response.session_id == "sess_123"
|
||||
assert response.retry_after == 5 # Running jobs suggest 5s retry
|
||||
|
||||
def test_job_to_result_response(self):
|
||||
"""Test to_result_response method."""
|
||||
job = Job(prompt="Test")
|
||||
job.start()
|
||||
job.succeed("Final result", {"tokens": 200})
|
||||
|
||||
response = job.to_result_response()
|
||||
|
||||
assert isinstance(response, JobResultResponse)
|
||||
assert response.job_id == job.id
|
||||
assert response.status == JobStatus.SUCCEEDED
|
||||
assert response.result == "Final result"
|
||||
assert response.metrics == {"tokens": 200}
|
||||
assert response.duration_seconds is not None
|
||||
|
||||
|
||||
class TestJobSubmitRequest:
|
||||
"""Tests for JobSubmitRequest model."""
|
||||
|
||||
def test_submit_request_minimal(self):
|
||||
"""Test minimal submit request."""
|
||||
request = JobSubmitRequest(prompt="Do something")
|
||||
|
||||
assert request.prompt == "Do something"
|
||||
assert request.framework == "praisonai"
|
||||
assert request.timeout == 3600
|
||||
|
||||
def test_submit_request_full(self):
|
||||
"""Test full submit request."""
|
||||
request = JobSubmitRequest(
|
||||
prompt="Do something",
|
||||
agent_file="agents.yaml",
|
||||
framework="crewai",
|
||||
webhook_url="https://example.com/hook",
|
||||
timeout=1800,
|
||||
session_id="sess_abc",
|
||||
idempotency_key="idem_xyz"
|
||||
)
|
||||
|
||||
assert request.agent_file == "agents.yaml"
|
||||
assert request.framework == "crewai"
|
||||
assert request.webhook_url == "https://example.com/hook"
|
||||
assert request.timeout == 1800
|
||||
assert request.session_id == "sess_abc"
|
||||
assert request.idempotency_key == "idem_xyz"
|
||||
|
||||
|
||||
class TestJobProgress:
|
||||
"""Tests for JobProgress model."""
|
||||
|
||||
def test_progress_defaults(self):
|
||||
"""Test progress defaults."""
|
||||
progress = JobProgress()
|
||||
|
||||
assert progress.percentage == 0.0
|
||||
assert progress.current_step is None
|
||||
assert progress.steps_completed == 0
|
||||
assert progress.steps_total is None
|
||||
|
||||
def test_progress_with_values(self):
|
||||
"""Test progress with values."""
|
||||
progress = JobProgress(
|
||||
percentage=75.0,
|
||||
current_step="Analyzing",
|
||||
steps_completed=3,
|
||||
steps_total=4
|
||||
)
|
||||
|
||||
assert progress.percentage == 75.0
|
||||
assert progress.current_step == "Analyzing"
|
||||
assert progress.steps_completed == 3
|
||||
assert progress.steps_total == 4
|
||||
|
||||
|
||||
class TestJobListResponse:
|
||||
"""Tests for JobListResponse model."""
|
||||
|
||||
def test_list_response(self):
|
||||
"""Test list response."""
|
||||
job1 = Job(prompt="Test 1")
|
||||
job2 = Job(prompt="Test 2")
|
||||
|
||||
response = JobListResponse(
|
||||
jobs=[job1.to_status_response(), job2.to_status_response()],
|
||||
total=2,
|
||||
page=1,
|
||||
page_size=20
|
||||
)
|
||||
|
||||
assert len(response.jobs) == 2
|
||||
assert response.total == 2
|
||||
assert response.page == 1
|
||||
assert response.page_size == 20
|
||||
@@ -0,0 +1,292 @@
|
||||
"""Integration tests for jobs router."""
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from praisonai.jobs.server import create_app
|
||||
from praisonai.jobs.store import InMemoryJobStore
|
||||
from praisonai.jobs.executor import JobExecutor
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store():
|
||||
"""Create a fresh store for each test."""
|
||||
return InMemoryJobStore(max_jobs=100)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor(store):
|
||||
"""Create executor with the store."""
|
||||
return JobExecutor(store=store, max_concurrent=5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(store, executor):
|
||||
"""Create test client."""
|
||||
app = create_app(store=store, executor=executor)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
class TestHealthEndpoint:
|
||||
"""Tests for health endpoint."""
|
||||
|
||||
def test_health_check(self, client):
|
||||
"""Test health check endpoint."""
|
||||
response = client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "healthy"
|
||||
assert "store" in data
|
||||
assert "executor_stats" in data
|
||||
|
||||
|
||||
class TestStatsEndpoint:
|
||||
"""Tests for stats endpoint."""
|
||||
|
||||
def test_get_stats(self, client):
|
||||
"""Test stats endpoint."""
|
||||
response = client.get("/stats")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "executor" in data
|
||||
|
||||
|
||||
class TestSubmitEndpoint:
|
||||
"""Tests for job submission."""
|
||||
|
||||
def test_submit_job(self, client):
|
||||
"""Test submitting a job."""
|
||||
response = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Test prompt"}
|
||||
)
|
||||
|
||||
assert response.status_code == 202
|
||||
data = response.json()
|
||||
assert "job_id" in data
|
||||
assert data["job_id"].startswith("run_")
|
||||
assert data["status"] == "queued"
|
||||
assert "poll_url" in data
|
||||
assert "stream_url" in data
|
||||
|
||||
def test_submit_job_with_all_fields(self, client):
|
||||
"""Test submitting a job with all fields."""
|
||||
response = client.post(
|
||||
"/api/v1/runs",
|
||||
json={
|
||||
"prompt": "Test prompt",
|
||||
"agent_file": "agents.yaml",
|
||||
"framework": "crewai",
|
||||
"timeout": 1800,
|
||||
"session_id": "sess_123"
|
||||
}
|
||||
)
|
||||
|
||||
assert response.status_code == 202
|
||||
data = response.json()
|
||||
assert data["status"] == "queued"
|
||||
|
||||
def test_submit_job_idempotency(self, client):
|
||||
"""Test idempotency key prevents duplicate jobs."""
|
||||
# First submission
|
||||
response1 = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Test prompt"},
|
||||
headers={"Idempotency-Key": "idem_test_123"}
|
||||
)
|
||||
assert response1.status_code == 202
|
||||
job_id1 = response1.json()["job_id"]
|
||||
|
||||
# Second submission with same key
|
||||
response2 = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Different prompt"},
|
||||
headers={"Idempotency-Key": "idem_test_123"}
|
||||
)
|
||||
assert response2.status_code == 202
|
||||
job_id2 = response2.json()["job_id"]
|
||||
|
||||
# Should return the same job
|
||||
assert job_id1 == job_id2
|
||||
|
||||
|
||||
class TestStatusEndpoint:
|
||||
"""Tests for job status."""
|
||||
|
||||
def test_get_status(self, client):
|
||||
"""Test getting job status."""
|
||||
# Submit a job first
|
||||
submit_response = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Test prompt"}
|
||||
)
|
||||
job_id = submit_response.json()["job_id"]
|
||||
|
||||
# Get status
|
||||
response = client.get(f"/api/v1/runs/{job_id}")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["job_id"] == job_id
|
||||
assert "status" in data
|
||||
assert "progress" in data
|
||||
|
||||
def test_get_status_not_found(self, client):
|
||||
"""Test getting status of nonexistent job."""
|
||||
response = client.get("/api/v1/runs/nonexistent_id")
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
class TestListEndpoint:
|
||||
"""Tests for listing jobs."""
|
||||
|
||||
def test_list_jobs_empty(self, client):
|
||||
"""Test listing jobs when empty."""
|
||||
response = client.get("/api/v1/runs")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["jobs"] == []
|
||||
assert data["total"] == 0
|
||||
|
||||
def test_list_jobs(self, client):
|
||||
"""Test listing jobs."""
|
||||
# Submit some jobs
|
||||
client.post("/api/v1/runs", json={"prompt": "Test 1"})
|
||||
client.post("/api/v1/runs", json={"prompt": "Test 2"})
|
||||
|
||||
response = client.get("/api/v1/runs")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data["jobs"]) == 2
|
||||
assert data["total"] == 2
|
||||
|
||||
def test_list_jobs_with_status_filter(self, client):
|
||||
"""Test listing jobs with status filter."""
|
||||
# Submit a job
|
||||
client.post("/api/v1/runs", json={"prompt": "Test"})
|
||||
|
||||
# Jobs may transition quickly, just verify the endpoint works
|
||||
response = client.get("/api/v1/runs?status=queued")
|
||||
assert response.status_code == 200
|
||||
|
||||
# Also test with other statuses
|
||||
response = client.get("/api/v1/runs?status=running")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = client.get("/api/v1/runs?status=failed")
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_list_jobs_pagination(self, client):
|
||||
"""Test listing jobs with pagination."""
|
||||
# Submit 5 jobs
|
||||
for i in range(5):
|
||||
client.post("/api/v1/runs", json={"prompt": f"Test {i}"})
|
||||
|
||||
# Get first page
|
||||
response = client.get("/api/v1/runs?page=1&page_size=2")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data["jobs"]) == 2
|
||||
assert data["total"] == 5
|
||||
assert data["page"] == 1
|
||||
assert data["page_size"] == 2
|
||||
|
||||
|
||||
class TestResultEndpoint:
|
||||
"""Tests for job results."""
|
||||
|
||||
def test_get_result_job_lifecycle(self, client):
|
||||
"""Test getting result follows job lifecycle."""
|
||||
# Submit a job
|
||||
submit_response = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Test prompt"}
|
||||
)
|
||||
job_id = submit_response.json()["job_id"]
|
||||
|
||||
# Result endpoint should work (may return 409 if not complete, or 200 if complete/failed)
|
||||
response = client.get(f"/api/v1/runs/{job_id}/result")
|
||||
|
||||
# Either 200 (completed) or 409 (not complete) is valid
|
||||
assert response.status_code in (200, 409)
|
||||
|
||||
def test_get_result_not_found(self, client):
|
||||
"""Test getting result of nonexistent job."""
|
||||
response = client.get("/api/v1/runs/nonexistent_id/result")
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
class TestCancelEndpoint:
|
||||
"""Tests for job cancellation."""
|
||||
|
||||
def test_cancel_job(self, client):
|
||||
"""Test cancelling a job."""
|
||||
# Submit a job
|
||||
submit_response = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Test prompt"}
|
||||
)
|
||||
job_id = submit_response.json()["job_id"]
|
||||
|
||||
# Cancel the job - may succeed (200) or fail if job already completed (409)
|
||||
response = client.post(f"/api/v1/runs/{job_id}/cancel")
|
||||
|
||||
# Either 200 (cancelled) or 409 (already complete) is valid due to async execution
|
||||
assert response.status_code in (200, 409)
|
||||
|
||||
def test_cancel_job_not_found(self, client):
|
||||
"""Test cancelling nonexistent job."""
|
||||
response = client.post("/api/v1/runs/nonexistent_id/cancel")
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
class TestDeleteEndpoint:
|
||||
"""Tests for job deletion."""
|
||||
|
||||
def test_delete_job_lifecycle(self, client):
|
||||
"""Test deleting a job follows lifecycle rules."""
|
||||
# Submit a job
|
||||
submit_response = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Test prompt"}
|
||||
)
|
||||
job_id = submit_response.json()["job_id"]
|
||||
|
||||
# Try to delete - may fail (409) if running, or succeed (204) if already complete
|
||||
response = client.delete(f"/api/v1/runs/{job_id}")
|
||||
|
||||
# Either 204 (deleted) or 409 (still running) is valid
|
||||
assert response.status_code in (204, 409)
|
||||
|
||||
def test_delete_cancelled_job(self, client):
|
||||
"""Test deleting a cancelled job."""
|
||||
# Submit and cancel a job
|
||||
submit_response = client.post(
|
||||
"/api/v1/runs",
|
||||
json={"prompt": "Test prompt"}
|
||||
)
|
||||
job_id = submit_response.json()["job_id"]
|
||||
client.post(f"/api/v1/runs/{job_id}/cancel")
|
||||
|
||||
# Delete the job
|
||||
response = client.delete(f"/api/v1/runs/{job_id}")
|
||||
|
||||
assert response.status_code == 204
|
||||
|
||||
# Verify it's gone
|
||||
get_response = client.get(f"/api/v1/runs/{job_id}")
|
||||
assert get_response.status_code == 404
|
||||
|
||||
def test_delete_not_found(self, client):
|
||||
"""Test deleting nonexistent job."""
|
||||
response = client.delete("/api/v1/runs/nonexistent_id")
|
||||
|
||||
assert response.status_code == 404
|
||||
@@ -0,0 +1,239 @@
|
||||
"""Unit tests for jobs store."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonai.jobs.models import Job, JobStatus
|
||||
from praisonai.jobs.store import InMemoryJobStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store():
|
||||
"""Create a fresh store for each test."""
|
||||
return InMemoryJobStore(max_jobs=100)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_job():
|
||||
"""Create a sample job."""
|
||||
return Job(prompt="Test prompt", session_id="sess_123")
|
||||
|
||||
|
||||
class TestInMemoryJobStore:
|
||||
"""Tests for InMemoryJobStore."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_and_get(self, store, sample_job):
|
||||
"""Test saving and retrieving a job."""
|
||||
await store.save(sample_job)
|
||||
|
||||
retrieved = await store.get(sample_job.id)
|
||||
|
||||
assert retrieved is not None
|
||||
assert retrieved.id == sample_job.id
|
||||
assert retrieved.prompt == "Test prompt"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nonexistent(self, store):
|
||||
"""Test getting a nonexistent job."""
|
||||
result = await store.get("nonexistent_id")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_idempotency_key(self, store):
|
||||
"""Test getting job by idempotency key."""
|
||||
job = Job(prompt="Test", idempotency_key="idem_123")
|
||||
await store.save(job)
|
||||
|
||||
retrieved = await store.get_by_idempotency_key("idem_123")
|
||||
|
||||
assert retrieved is not None
|
||||
assert retrieved.id == job.id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_idempotency_key_nonexistent(self, store):
|
||||
"""Test getting by nonexistent idempotency key."""
|
||||
result = await store.get_by_idempotency_key("nonexistent")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_jobs_empty(self, store):
|
||||
"""Test listing jobs when empty."""
|
||||
jobs = await store.list_jobs()
|
||||
assert jobs == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_jobs(self, store):
|
||||
"""Test listing jobs."""
|
||||
job1 = Job(prompt="Test 1")
|
||||
job2 = Job(prompt="Test 2")
|
||||
await store.save(job1)
|
||||
await store.save(job2)
|
||||
|
||||
jobs = await store.list_jobs()
|
||||
|
||||
assert len(jobs) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_jobs_by_status(self, store):
|
||||
"""Test listing jobs filtered by status."""
|
||||
job1 = Job(prompt="Test 1")
|
||||
job2 = Job(prompt="Test 2")
|
||||
job2.start()
|
||||
await store.save(job1)
|
||||
await store.save(job2)
|
||||
|
||||
queued = await store.list_jobs(status=JobStatus.QUEUED)
|
||||
running = await store.list_jobs(status=JobStatus.RUNNING)
|
||||
|
||||
assert len(queued) == 1
|
||||
assert len(running) == 1
|
||||
assert queued[0].id == job1.id
|
||||
assert running[0].id == job2.id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_jobs_by_session_id(self, store):
|
||||
"""Test listing jobs filtered by session ID."""
|
||||
job1 = Job(prompt="Test 1", session_id="sess_a")
|
||||
job2 = Job(prompt="Test 2", session_id="sess_b")
|
||||
job3 = Job(prompt="Test 3", session_id="sess_a")
|
||||
await store.save(job1)
|
||||
await store.save(job2)
|
||||
await store.save(job3)
|
||||
|
||||
sess_a_jobs = await store.list_jobs(session_id="sess_a")
|
||||
|
||||
assert len(sess_a_jobs) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_jobs_pagination(self, store):
|
||||
"""Test listing jobs with pagination."""
|
||||
for i in range(5):
|
||||
await store.save(Job(prompt=f"Test {i}"))
|
||||
|
||||
page1 = await store.list_jobs(limit=2, offset=0)
|
||||
page2 = await store.list_jobs(limit=2, offset=2)
|
||||
page3 = await store.list_jobs(limit=2, offset=4)
|
||||
|
||||
assert len(page1) == 2
|
||||
assert len(page2) == 2
|
||||
assert len(page3) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count(self, store):
|
||||
"""Test counting jobs."""
|
||||
for i in range(3):
|
||||
await store.save(Job(prompt=f"Test {i}"))
|
||||
|
||||
count = await store.count()
|
||||
assert count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_by_status(self, store):
|
||||
"""Test counting jobs by status."""
|
||||
job1 = Job(prompt="Test 1")
|
||||
job2 = Job(prompt="Test 2")
|
||||
job2.start()
|
||||
job3 = Job(prompt="Test 3")
|
||||
job3.start()
|
||||
job3.succeed("Done")
|
||||
|
||||
await store.save(job1)
|
||||
await store.save(job2)
|
||||
await store.save(job3)
|
||||
|
||||
queued_count = await store.count(status=JobStatus.QUEUED)
|
||||
running_count = await store.count(status=JobStatus.RUNNING)
|
||||
succeeded_count = await store.count(status=JobStatus.SUCCEEDED)
|
||||
|
||||
assert queued_count == 1
|
||||
assert running_count == 1
|
||||
assert succeeded_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete(self, store, sample_job):
|
||||
"""Test deleting a job."""
|
||||
await store.save(sample_job)
|
||||
|
||||
result = await store.delete(sample_job.id)
|
||||
|
||||
assert result is True
|
||||
assert await store.get(sample_job.id) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_nonexistent(self, store):
|
||||
"""Test deleting a nonexistent job."""
|
||||
result = await store.delete("nonexistent")
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_removes_idempotency_key(self, store):
|
||||
"""Test that deleting a job removes its idempotency key."""
|
||||
job = Job(prompt="Test", idempotency_key="idem_456")
|
||||
await store.save(job)
|
||||
|
||||
await store.delete(job.id)
|
||||
|
||||
assert await store.get_by_idempotency_key("idem_456") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_old_jobs(self, store):
|
||||
"""Test cleaning up old completed jobs."""
|
||||
# Create completed jobs
|
||||
old_job = Job(prompt="Old")
|
||||
old_job.start()
|
||||
old_job.succeed("Done")
|
||||
# Manually set completed_at to past
|
||||
from datetime import datetime, timedelta
|
||||
old_job.completed_at = datetime.utcnow() - timedelta(days=2)
|
||||
|
||||
new_job = Job(prompt="New")
|
||||
new_job.start()
|
||||
new_job.succeed("Done")
|
||||
|
||||
await store.save(old_job)
|
||||
await store.save(new_job)
|
||||
|
||||
# Cleanup jobs older than 1 day
|
||||
cleaned = await store.cleanup_old_jobs(max_age_seconds=86400)
|
||||
|
||||
assert cleaned == 1
|
||||
assert await store.get(old_job.id) is None
|
||||
assert await store.get(new_job.id) is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_jobs_eviction(self):
|
||||
"""Test that old completed jobs are evicted when max is reached."""
|
||||
store = InMemoryJobStore(max_jobs=5)
|
||||
|
||||
# Create 5 completed jobs
|
||||
for i in range(5):
|
||||
job = Job(prompt=f"Test {i}")
|
||||
job.start()
|
||||
job.succeed(f"Result {i}")
|
||||
await store.save(job)
|
||||
|
||||
# Add one more - should trigger eviction
|
||||
new_job = Job(prompt="New job")
|
||||
await store.save(new_job)
|
||||
|
||||
# Should still have at most max_jobs
|
||||
count = await store.count()
|
||||
assert count <= 6 # May have evicted some
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stats(self, store):
|
||||
"""Test getting store statistics."""
|
||||
job1 = Job(prompt="Test 1")
|
||||
job2 = Job(prompt="Test 2", idempotency_key="idem_1")
|
||||
job2.start()
|
||||
await store.save(job1)
|
||||
await store.save(job2)
|
||||
|
||||
stats = store.get_stats()
|
||||
|
||||
assert stats["total_jobs"] == 2
|
||||
assert stats["idempotency_keys"] == 1
|
||||
assert stats["max_jobs"] == 100
|
||||
assert "status_counts" in stats
|
||||
assert stats["status_counts"]["queued"] == 1
|
||||
assert stats["status_counts"]["running"] == 1
|
||||
@@ -0,0 +1 @@
|
||||
"""Unit tests for knowledge stack."""
|
||||
@@ -0,0 +1,224 @@
|
||||
"""
|
||||
Unit tests for Index Type Selection.
|
||||
"""
|
||||
|
||||
from praisonaiagents.knowledge.index import (
|
||||
IndexType,
|
||||
IndexStats,
|
||||
get_index_registry,
|
||||
KeywordIndex,
|
||||
)
|
||||
|
||||
|
||||
class TestIndexType:
|
||||
"""Tests for IndexType enum."""
|
||||
|
||||
def test_index_type_values(self):
|
||||
"""Test index type enum values."""
|
||||
assert IndexType.VECTOR == "vector"
|
||||
assert IndexType.KEYWORD == "keyword"
|
||||
assert IndexType.HYBRID == "hybrid"
|
||||
assert IndexType.GRAPH == "graph"
|
||||
|
||||
def test_index_type_from_string(self):
|
||||
"""Test creating index type from string."""
|
||||
assert IndexType("vector") == IndexType.VECTOR
|
||||
assert IndexType("keyword") == IndexType.KEYWORD
|
||||
|
||||
|
||||
class TestIndexStats:
|
||||
"""Tests for IndexStats dataclass."""
|
||||
|
||||
def test_stats_creation(self):
|
||||
"""Test basic stats creation."""
|
||||
stats = IndexStats(
|
||||
index_type=IndexType.VECTOR,
|
||||
document_count=100,
|
||||
chunk_count=500
|
||||
)
|
||||
assert stats.index_type == IndexType.VECTOR
|
||||
assert stats.document_count == 100
|
||||
assert stats.chunk_count == 500
|
||||
assert stats.total_tokens == 0
|
||||
|
||||
def test_stats_with_metadata(self):
|
||||
"""Test stats with metadata."""
|
||||
stats = IndexStats(
|
||||
index_type=IndexType.KEYWORD,
|
||||
document_count=50,
|
||||
chunk_count=200,
|
||||
total_tokens=10000,
|
||||
metadata={"unique_terms": 5000}
|
||||
)
|
||||
assert stats.total_tokens == 10000
|
||||
assert stats.metadata["unique_terms"] == 5000
|
||||
|
||||
|
||||
class TestKeywordIndex:
|
||||
"""Tests for KeywordIndex (BM25)."""
|
||||
|
||||
def test_add_documents(self):
|
||||
"""Test adding documents."""
|
||||
index = KeywordIndex()
|
||||
ids = index.add_documents(
|
||||
texts=["Hello world", "Goodbye world"],
|
||||
)
|
||||
assert len(ids) == 2
|
||||
|
||||
stats = index.stats()
|
||||
assert stats.document_count == 2
|
||||
|
||||
def test_add_with_metadata(self):
|
||||
"""Test adding documents with metadata."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(
|
||||
texts=["Test document"],
|
||||
metadatas=[{"source": "test.txt"}]
|
||||
)
|
||||
|
||||
results = index.query("Test")
|
||||
assert len(results) > 0
|
||||
assert results[0]["metadata"]["source"] == "test.txt"
|
||||
|
||||
def test_query_basic(self):
|
||||
"""Test basic querying."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(
|
||||
texts=[
|
||||
"The quick brown fox jumps over the lazy dog",
|
||||
"Python is a programming language",
|
||||
"Machine learning with neural networks"
|
||||
]
|
||||
)
|
||||
|
||||
results = index.query("fox dog")
|
||||
|
||||
assert len(results) > 0
|
||||
# First result should be about fox and dog
|
||||
assert "fox" in results[0]["text"].lower() or "dog" in results[0]["text"].lower()
|
||||
|
||||
def test_query_with_top_k(self):
|
||||
"""Test querying with top_k limit."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(
|
||||
texts=["Doc A", "Doc B", "Doc C", "Doc D"]
|
||||
)
|
||||
|
||||
results = index.query("Doc", top_k=2)
|
||||
assert len(results) == 2
|
||||
|
||||
def test_query_with_filter(self):
|
||||
"""Test querying with metadata filter."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(
|
||||
texts=["Doc A", "Doc B", "Doc C"],
|
||||
metadatas=[
|
||||
{"type": "a"},
|
||||
{"type": "b"},
|
||||
{"type": "a"}
|
||||
]
|
||||
)
|
||||
|
||||
results = index.query("Doc", filter={"type": "a"})
|
||||
|
||||
assert len(results) == 2
|
||||
for r in results:
|
||||
assert r["metadata"]["type"] == "a"
|
||||
|
||||
def test_delete_by_ids(self):
|
||||
"""Test deleting by IDs."""
|
||||
index = KeywordIndex()
|
||||
ids = index.add_documents(
|
||||
texts=["A", "B", "C"]
|
||||
)
|
||||
|
||||
deleted = index.delete(ids=[ids[0], ids[1]])
|
||||
assert deleted == 2
|
||||
|
||||
stats = index.stats()
|
||||
assert stats.document_count == 1
|
||||
|
||||
def test_delete_all(self):
|
||||
"""Test deleting all documents."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(texts=["A", "B", "C"])
|
||||
|
||||
deleted = index.delete(delete_all=True)
|
||||
assert deleted == 3
|
||||
|
||||
stats = index.stats()
|
||||
assert stats.document_count == 0
|
||||
|
||||
def test_delete_by_filter(self):
|
||||
"""Test deleting by filter."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(
|
||||
texts=["A", "B", "C"],
|
||||
metadatas=[{"keep": True}, {"keep": False}, {"keep": True}]
|
||||
)
|
||||
|
||||
deleted = index.delete(filter={"keep": False})
|
||||
assert deleted == 1
|
||||
|
||||
stats = index.stats()
|
||||
assert stats.document_count == 2
|
||||
|
||||
def test_stats(self):
|
||||
"""Test getting index statistics."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(
|
||||
texts=["Hello world", "Test document"]
|
||||
)
|
||||
|
||||
stats = index.stats()
|
||||
|
||||
assert stats.index_type == IndexType.KEYWORD
|
||||
assert stats.document_count == 2
|
||||
assert stats.chunk_count == 2
|
||||
assert stats.total_tokens > 0
|
||||
assert "unique_terms" in stats.metadata
|
||||
|
||||
def test_bm25_scoring(self):
|
||||
"""Test BM25 scoring behavior."""
|
||||
index = KeywordIndex()
|
||||
index.add_documents(
|
||||
texts=[
|
||||
"apple apple apple", # High TF for apple
|
||||
"apple banana", # Lower TF for apple
|
||||
"banana banana" # No apple
|
||||
]
|
||||
)
|
||||
|
||||
results = index.query("apple")
|
||||
|
||||
# First result should have highest score (most apples)
|
||||
assert results[0]["text"] == "apple apple apple"
|
||||
# Last result should have lowest score (no apples)
|
||||
assert results[-1]["text"] == "banana banana"
|
||||
|
||||
|
||||
class TestIndexRegistry:
|
||||
"""Tests for IndexRegistry."""
|
||||
|
||||
def test_keyword_index_registered(self):
|
||||
"""Test that keyword index is registered by default."""
|
||||
registry = get_index_registry()
|
||||
assert "keyword" in registry.list_indices()
|
||||
|
||||
def test_get_keyword_index(self):
|
||||
"""Test getting keyword index."""
|
||||
registry = get_index_registry()
|
||||
index = registry.get("keyword")
|
||||
assert index is not None
|
||||
assert index.index_type == IndexType.KEYWORD
|
||||
|
||||
def test_get_nonexistent_index(self):
|
||||
"""Test getting non-existent index."""
|
||||
registry = get_index_registry()
|
||||
assert registry.get("nonexistent") is None
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Test that registry is a singleton."""
|
||||
registry1 = get_index_registry()
|
||||
registry2 = get_index_registry()
|
||||
assert registry1 is registry2
|
||||
@@ -0,0 +1,250 @@
|
||||
"""
|
||||
Unit tests for Query Engine Patterns.
|
||||
"""
|
||||
|
||||
from praisonaiagents.knowledge.query_engine import (
|
||||
QueryMode,
|
||||
QueryResult,
|
||||
get_query_engine_registry,
|
||||
decompose_question,
|
||||
synthesize_answer,
|
||||
SimpleQueryEngine,
|
||||
SubQuestionEngine,
|
||||
)
|
||||
|
||||
|
||||
class TestQueryMode:
|
||||
"""Tests for QueryMode enum."""
|
||||
|
||||
def test_query_mode_values(self):
|
||||
"""Test query mode enum values."""
|
||||
assert QueryMode.DEFAULT == "default"
|
||||
assert QueryMode.SUB_QUESTION == "sub_question"
|
||||
assert QueryMode.SQL == "sql"
|
||||
assert QueryMode.ROUTER == "router"
|
||||
assert QueryMode.SUMMARIZE == "summarize"
|
||||
|
||||
def test_query_mode_from_string(self):
|
||||
"""Test creating query mode from string."""
|
||||
assert QueryMode("default") == QueryMode.DEFAULT
|
||||
assert QueryMode("sub_question") == QueryMode.SUB_QUESTION
|
||||
|
||||
|
||||
class TestQueryResult:
|
||||
"""Tests for QueryResult dataclass."""
|
||||
|
||||
def test_result_creation(self):
|
||||
"""Test basic result creation."""
|
||||
result = QueryResult(answer="The answer is 42")
|
||||
assert result.answer == "The answer is 42"
|
||||
assert result.sources == []
|
||||
assert result.sub_questions is None
|
||||
assert result.metadata == {}
|
||||
|
||||
def test_result_with_sources(self):
|
||||
"""Test result with sources."""
|
||||
result = QueryResult(
|
||||
answer="Test answer",
|
||||
sources=[
|
||||
{"text": "Source 1", "score": 0.9},
|
||||
{"text": "Source 2", "score": 0.8}
|
||||
]
|
||||
)
|
||||
assert len(result.sources) == 2
|
||||
assert result.sources[0]["score"] == 0.9
|
||||
|
||||
def test_result_with_sub_questions(self):
|
||||
"""Test result with sub-questions."""
|
||||
result = QueryResult(
|
||||
answer="Combined answer",
|
||||
sub_questions=["What is X?", "What is Y?"]
|
||||
)
|
||||
assert len(result.sub_questions) == 2
|
||||
|
||||
def test_result_to_dict(self):
|
||||
"""Test result serialization."""
|
||||
result = QueryResult(
|
||||
answer="Test",
|
||||
sources=[{"text": "Source"}],
|
||||
metadata={"mode": "default"}
|
||||
)
|
||||
d = result.to_dict()
|
||||
assert d["answer"] == "Test"
|
||||
assert len(d["sources"]) == 1
|
||||
assert d["metadata"]["mode"] == "default"
|
||||
|
||||
|
||||
class TestDecomposeQuestion:
|
||||
"""Tests for question decomposition."""
|
||||
|
||||
def test_simple_question(self):
|
||||
"""Test that simple questions are not decomposed."""
|
||||
questions = decompose_question("What is Python?")
|
||||
assert len(questions) == 1
|
||||
assert questions[0] == "What is Python?"
|
||||
|
||||
def test_compound_question_and(self):
|
||||
"""Test decomposing compound question with 'and'."""
|
||||
questions = decompose_question("What is Python and what is JavaScript?")
|
||||
assert len(questions) >= 2
|
||||
|
||||
def test_compound_question_also(self):
|
||||
"""Test decomposing compound question with 'also'."""
|
||||
questions = decompose_question("Explain Python also explain JavaScript")
|
||||
assert len(questions) >= 2
|
||||
|
||||
def test_question_mark_added(self):
|
||||
"""Test that question marks are added."""
|
||||
questions = decompose_question("What is X and what is Y")
|
||||
for q in questions:
|
||||
assert q.endswith("?")
|
||||
|
||||
def test_capitalization(self):
|
||||
"""Test that questions are capitalized."""
|
||||
questions = decompose_question("what is x and what is y")
|
||||
for q in questions:
|
||||
assert q[0].isupper()
|
||||
|
||||
|
||||
class TestSynthesizeAnswer:
|
||||
"""Tests for answer synthesis."""
|
||||
|
||||
def test_empty_contexts(self):
|
||||
"""Test synthesis with no contexts."""
|
||||
answer = synthesize_answer("What is X?", [])
|
||||
assert "No relevant information" in answer
|
||||
|
||||
def test_single_context(self):
|
||||
"""Test synthesis with single context."""
|
||||
answer = synthesize_answer(
|
||||
"What is Python?",
|
||||
["Python is a programming language."]
|
||||
)
|
||||
assert "Python is a programming language" in answer
|
||||
|
||||
def test_multiple_contexts(self):
|
||||
"""Test synthesis with multiple contexts."""
|
||||
answer = synthesize_answer(
|
||||
"What is Python?",
|
||||
[
|
||||
"Python is a programming language.",
|
||||
"Python was created by Guido van Rossum."
|
||||
]
|
||||
)
|
||||
assert "Python" in answer
|
||||
|
||||
def test_context_truncation(self):
|
||||
"""Test that long contexts are truncated."""
|
||||
long_context = "x" * 5000
|
||||
answer = synthesize_answer("Question?", [long_context], max_context_length=100)
|
||||
# Answer should be shorter than the original context
|
||||
assert len(answer) < 5000
|
||||
|
||||
|
||||
class TestSimpleQueryEngine:
|
||||
"""Tests for SimpleQueryEngine."""
|
||||
|
||||
def test_query_with_context(self):
|
||||
"""Test querying with context."""
|
||||
engine = SimpleQueryEngine()
|
||||
result = engine.query(
|
||||
"What is Python?",
|
||||
context=["Python is a programming language."]
|
||||
)
|
||||
|
||||
assert result.answer is not None
|
||||
assert len(result.sources) == 1
|
||||
assert result.metadata["mode"] == "default"
|
||||
|
||||
def test_query_without_context(self):
|
||||
"""Test querying without context."""
|
||||
engine = SimpleQueryEngine()
|
||||
result = engine.query("What is Python?")
|
||||
|
||||
assert "No relevant information" in result.answer
|
||||
assert len(result.sources) == 0
|
||||
|
||||
def test_aquery_async(self):
|
||||
"""Test async querying."""
|
||||
import asyncio
|
||||
|
||||
engine = SimpleQueryEngine()
|
||||
|
||||
async def run_test():
|
||||
result = await engine.aquery(
|
||||
"What is Python?",
|
||||
context=["Python is a language."]
|
||||
)
|
||||
return result
|
||||
|
||||
result = asyncio.run(run_test())
|
||||
assert result.answer is not None
|
||||
|
||||
|
||||
class TestSubQuestionEngine:
|
||||
"""Tests for SubQuestionEngine."""
|
||||
|
||||
def test_query_decomposes(self):
|
||||
"""Test that query decomposes questions."""
|
||||
engine = SubQuestionEngine()
|
||||
result = engine.query(
|
||||
"What is Python and what is JavaScript?",
|
||||
context=["Python is a language.", "JavaScript is also a language."]
|
||||
)
|
||||
|
||||
assert result.answer is not None
|
||||
assert result.sub_questions is not None
|
||||
assert len(result.sub_questions) >= 1
|
||||
|
||||
def test_query_simple_question(self):
|
||||
"""Test with simple question (no decomposition needed)."""
|
||||
engine = SubQuestionEngine()
|
||||
result = engine.query(
|
||||
"What is Python?",
|
||||
context=["Python is a programming language."]
|
||||
)
|
||||
|
||||
assert result.answer is not None
|
||||
|
||||
def test_metadata_includes_count(self):
|
||||
"""Test that metadata includes sub-question count."""
|
||||
engine = SubQuestionEngine()
|
||||
result = engine.query(
|
||||
"What is X and what is Y?",
|
||||
context=["X is something.", "Y is something else."]
|
||||
)
|
||||
|
||||
assert "sub_question_count" in result.metadata
|
||||
|
||||
|
||||
class TestQueryEngineRegistry:
|
||||
"""Tests for QueryEngineRegistry."""
|
||||
|
||||
def test_default_engine_registered(self):
|
||||
"""Test that default engine is registered."""
|
||||
registry = get_query_engine_registry()
|
||||
assert "default" in registry.list_engines()
|
||||
assert "simple" in registry.list_engines()
|
||||
|
||||
def test_sub_question_engine_registered(self):
|
||||
"""Test that sub_question engine is registered."""
|
||||
registry = get_query_engine_registry()
|
||||
assert "sub_question" in registry.list_engines()
|
||||
|
||||
def test_get_default_engine(self):
|
||||
"""Test getting default engine."""
|
||||
registry = get_query_engine_registry()
|
||||
engine = registry.get("default")
|
||||
assert engine is not None
|
||||
assert engine.mode == QueryMode.DEFAULT
|
||||
|
||||
def test_get_nonexistent_engine(self):
|
||||
"""Test getting non-existent engine."""
|
||||
registry = get_query_engine_registry()
|
||||
assert registry.get("nonexistent") is None
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Test that registry is a singleton."""
|
||||
registry1 = get_query_engine_registry()
|
||||
registry2 = get_query_engine_registry()
|
||||
assert registry1 is registry2
|
||||
@@ -0,0 +1,220 @@
|
||||
"""
|
||||
Unit tests for Reader Protocol and Registry.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from praisonaiagents.knowledge.readers import (
|
||||
Document,
|
||||
get_reader_registry,
|
||||
detect_source_kind,
|
||||
get_file_extension,
|
||||
)
|
||||
|
||||
|
||||
class TestDocument:
|
||||
"""Tests for Document dataclass."""
|
||||
|
||||
def test_document_creation(self):
|
||||
"""Test basic document creation."""
|
||||
doc = Document(content="Hello world")
|
||||
assert doc.content == "Hello world"
|
||||
assert doc.metadata == {}
|
||||
assert doc.doc_id is not None
|
||||
assert doc.embedding is None
|
||||
|
||||
def test_document_with_metadata(self):
|
||||
"""Test document with metadata."""
|
||||
doc = Document(
|
||||
content="Test content",
|
||||
metadata={"source": "test.txt", "page": 1}
|
||||
)
|
||||
assert doc.metadata["source"] == "test.txt"
|
||||
assert doc.metadata["page"] == 1
|
||||
|
||||
def test_document_to_dict(self):
|
||||
"""Test document serialization."""
|
||||
doc = Document(content="Test", metadata={"key": "value"})
|
||||
d = doc.to_dict()
|
||||
assert d["content"] == "Test"
|
||||
assert d["metadata"]["key"] == "value"
|
||||
assert "doc_id" in d
|
||||
|
||||
def test_document_from_dict(self):
|
||||
"""Test document deserialization."""
|
||||
data = {
|
||||
"content": "Test content",
|
||||
"metadata": {"source": "test.txt"},
|
||||
"doc_id": "test-id"
|
||||
}
|
||||
doc = Document.from_dict(data)
|
||||
assert doc.content == "Test content"
|
||||
assert doc.doc_id == "test-id"
|
||||
|
||||
|
||||
class TestDetectSourceKind:
|
||||
"""Tests for source kind detection."""
|
||||
|
||||
def test_detect_url_http(self):
|
||||
"""Test HTTP URL detection."""
|
||||
assert detect_source_kind("http://example.com") == "url"
|
||||
|
||||
def test_detect_url_https(self):
|
||||
"""Test HTTPS URL detection."""
|
||||
assert detect_source_kind("https://example.com/page.html") == "url"
|
||||
|
||||
def test_detect_url_s3(self):
|
||||
"""Test S3 URL detection."""
|
||||
assert detect_source_kind("s3://bucket/key") == "url"
|
||||
|
||||
def test_detect_glob_asterisk(self):
|
||||
"""Test glob pattern with asterisk."""
|
||||
assert detect_source_kind("*.txt") == "glob"
|
||||
assert detect_source_kind("docs/**/*.md") == "glob"
|
||||
|
||||
def test_detect_glob_question(self):
|
||||
"""Test glob pattern with question mark."""
|
||||
assert detect_source_kind("file?.txt") == "glob"
|
||||
|
||||
def test_detect_glob_bracket(self):
|
||||
"""Test glob pattern with brackets."""
|
||||
assert detect_source_kind("file[0-9].txt") == "glob"
|
||||
|
||||
def test_detect_file_with_extension(self):
|
||||
"""Test file path detection."""
|
||||
assert detect_source_kind("document.pdf") == "file"
|
||||
assert detect_source_kind("/path/to/file.txt") == "file"
|
||||
|
||||
def test_detect_directory(self):
|
||||
"""Test directory detection."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
assert detect_source_kind(tmpdir) == "directory"
|
||||
|
||||
def test_detect_existing_file(self):
|
||||
"""Test existing file detection."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
|
||||
f.write(b"test")
|
||||
path = f.name
|
||||
try:
|
||||
assert detect_source_kind(path) == "file"
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_detect_unknown(self):
|
||||
"""Test unknown source kind."""
|
||||
assert detect_source_kind("") == "unknown"
|
||||
assert detect_source_kind(" ") == "unknown"
|
||||
|
||||
|
||||
class TestGetFileExtension:
|
||||
"""Tests for file extension extraction."""
|
||||
|
||||
def test_simple_extension(self):
|
||||
"""Test simple file extension."""
|
||||
assert get_file_extension("file.txt") == "txt"
|
||||
assert get_file_extension("document.pdf") == "pdf"
|
||||
|
||||
def test_path_extension(self):
|
||||
"""Test extension from path."""
|
||||
assert get_file_extension("/path/to/file.docx") == "docx"
|
||||
|
||||
def test_url_extension(self):
|
||||
"""Test extension from URL."""
|
||||
assert get_file_extension("https://example.com/doc.pdf") == "pdf"
|
||||
|
||||
def test_no_extension(self):
|
||||
"""Test file without extension."""
|
||||
assert get_file_extension("README") == ""
|
||||
|
||||
def test_multiple_dots(self):
|
||||
"""Test file with multiple dots."""
|
||||
assert get_file_extension("file.tar.gz") == "gz"
|
||||
|
||||
|
||||
class TestReaderRegistry:
|
||||
"""Tests for ReaderRegistry."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Clear registry before each test."""
|
||||
registry = get_reader_registry()
|
||||
registry.clear()
|
||||
|
||||
def test_register_reader(self):
|
||||
"""Test registering a reader."""
|
||||
registry = get_reader_registry()
|
||||
|
||||
class MockReader:
|
||||
name = "mock"
|
||||
supported_extensions = ["mock"]
|
||||
def load(self, source, *, metadata=None):
|
||||
return [Document(content="mock")]
|
||||
def can_handle(self, source):
|
||||
return source.endswith(".mock")
|
||||
|
||||
registry.register("mock", MockReader, ["mock"])
|
||||
assert "mock" in registry.list_readers()
|
||||
|
||||
def test_get_reader(self):
|
||||
"""Test getting a reader by name."""
|
||||
registry = get_reader_registry()
|
||||
|
||||
class MockReader:
|
||||
name = "mock"
|
||||
supported_extensions = ["mock"]
|
||||
def load(self, source, *, metadata=None):
|
||||
return []
|
||||
def can_handle(self, source):
|
||||
return True
|
||||
|
||||
registry.register("mock", MockReader, ["mock"])
|
||||
reader = registry.get("mock")
|
||||
assert reader is not None
|
||||
assert reader.name == "mock"
|
||||
|
||||
def test_get_nonexistent_reader(self):
|
||||
"""Test getting a non-existent reader."""
|
||||
registry = get_reader_registry()
|
||||
assert registry.get("nonexistent") is None
|
||||
|
||||
def test_extension_mapping(self):
|
||||
"""Test extension to reader mapping."""
|
||||
registry = get_reader_registry()
|
||||
|
||||
class MockReader:
|
||||
name = "mock"
|
||||
supported_extensions = ["txt", "text"]
|
||||
def load(self, source, *, metadata=None):
|
||||
return []
|
||||
def can_handle(self, source):
|
||||
return True
|
||||
|
||||
registry.register("mock", MockReader, ["txt", "text"])
|
||||
|
||||
extensions = registry.list_extensions()
|
||||
assert "txt" in extensions
|
||||
assert "text" in extensions
|
||||
assert extensions["txt"] == "mock"
|
||||
|
||||
def test_get_for_source_file(self):
|
||||
"""Test getting reader for file source."""
|
||||
registry = get_reader_registry()
|
||||
|
||||
class TxtReader:
|
||||
name = "txt"
|
||||
supported_extensions = ["txt"]
|
||||
def load(self, source, *, metadata=None):
|
||||
return []
|
||||
def can_handle(self, source):
|
||||
return source.endswith(".txt")
|
||||
|
||||
registry.register("txt", TxtReader, ["txt"])
|
||||
|
||||
reader = registry.get_for_source("document.txt")
|
||||
assert reader is not None
|
||||
assert reader.name == "txt"
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Test that registry is a singleton."""
|
||||
registry1 = get_reader_registry()
|
||||
registry2 = get_reader_registry()
|
||||
assert registry1 is registry2
|
||||
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
Unit tests for Reranker Protocol and Registry.
|
||||
"""
|
||||
|
||||
from praisonaiagents.knowledge.rerankers import (
|
||||
RerankResult,
|
||||
get_reranker_registry,
|
||||
SimpleReranker,
|
||||
)
|
||||
|
||||
|
||||
class TestRerankResult:
|
||||
"""Tests for RerankResult dataclass."""
|
||||
|
||||
def test_result_creation(self):
|
||||
"""Test basic result creation."""
|
||||
result = RerankResult(text="Hello world", score=0.95, original_index=0)
|
||||
assert result.text == "Hello world"
|
||||
assert result.score == 0.95
|
||||
assert result.original_index == 0
|
||||
assert result.metadata == {}
|
||||
|
||||
def test_result_with_metadata(self):
|
||||
"""Test result with metadata."""
|
||||
result = RerankResult(
|
||||
text="Test content",
|
||||
score=0.8,
|
||||
original_index=1,
|
||||
metadata={"source": "test.txt"}
|
||||
)
|
||||
assert result.metadata["source"] == "test.txt"
|
||||
|
||||
|
||||
class TestSimpleReranker:
|
||||
"""Tests for SimpleReranker."""
|
||||
|
||||
def test_rerank_basic(self):
|
||||
"""Test basic reranking."""
|
||||
reranker = SimpleReranker()
|
||||
documents = [
|
||||
"The quick brown fox jumps over the lazy dog",
|
||||
"Python is a programming language",
|
||||
"Machine learning with neural networks"
|
||||
]
|
||||
|
||||
results = reranker.rerank("fox dog", documents)
|
||||
|
||||
assert len(results) == 3
|
||||
# First document should score highest (contains fox and dog)
|
||||
assert results[0].text == documents[0]
|
||||
|
||||
def test_rerank_with_top_k(self):
|
||||
"""Test reranking with top_k limit."""
|
||||
reranker = SimpleReranker()
|
||||
documents = ["Doc A", "Doc B", "Doc C", "Doc D"]
|
||||
|
||||
results = reranker.rerank("Doc", documents, top_k=2)
|
||||
|
||||
assert len(results) == 2
|
||||
|
||||
def test_rerank_preserves_original_index(self):
|
||||
"""Test that original index is preserved."""
|
||||
reranker = SimpleReranker()
|
||||
documents = ["First", "Second", "Third"]
|
||||
|
||||
results = reranker.rerank("Second", documents)
|
||||
|
||||
# Find the result for "Second"
|
||||
second_result = [r for r in results if r.text == "Second"][0]
|
||||
assert second_result.original_index == 1
|
||||
|
||||
def test_rerank_empty_documents(self):
|
||||
"""Test reranking with empty documents."""
|
||||
reranker = SimpleReranker()
|
||||
results = reranker.rerank("query", [])
|
||||
assert results == []
|
||||
|
||||
def test_rerank_empty_query(self):
|
||||
"""Test reranking with empty query."""
|
||||
reranker = SimpleReranker()
|
||||
documents = ["Doc A", "Doc B"]
|
||||
|
||||
results = reranker.rerank("", documents)
|
||||
|
||||
assert len(results) == 2
|
||||
# All should have score 0 with empty query
|
||||
for r in results:
|
||||
assert r.score == 0.0
|
||||
|
||||
def test_arerank_async(self):
|
||||
"""Test async reranking."""
|
||||
import asyncio
|
||||
|
||||
reranker = SimpleReranker()
|
||||
documents = ["Hello world", "Goodbye world"]
|
||||
|
||||
async def run_test():
|
||||
results = await reranker.arerank("Hello", documents)
|
||||
return results
|
||||
|
||||
results = asyncio.run(run_test())
|
||||
assert len(results) == 2
|
||||
|
||||
|
||||
class TestRerankerRegistry:
|
||||
"""Tests for RerankerRegistry."""
|
||||
|
||||
def test_simple_reranker_registered(self):
|
||||
"""Test that simple reranker is registered by default."""
|
||||
registry = get_reranker_registry()
|
||||
assert "simple" in registry.list_rerankers()
|
||||
|
||||
def test_get_simple_reranker(self):
|
||||
"""Test getting simple reranker."""
|
||||
registry = get_reranker_registry()
|
||||
reranker = registry.get("simple")
|
||||
assert reranker is not None
|
||||
assert reranker.name == "simple"
|
||||
|
||||
def test_get_nonexistent_reranker(self):
|
||||
"""Test getting non-existent reranker."""
|
||||
registry = get_reranker_registry()
|
||||
assert registry.get("nonexistent") is None
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Test that registry is a singleton."""
|
||||
registry1 = get_reranker_registry()
|
||||
registry2 = get_reranker_registry()
|
||||
assert registry1 is registry2
|
||||
@@ -0,0 +1,200 @@
|
||||
"""
|
||||
Unit tests for Retrieval Strategy Patterns.
|
||||
"""
|
||||
|
||||
from praisonaiagents.knowledge.retrieval import (
|
||||
RetrievalResult,
|
||||
RetrievalStrategy,
|
||||
get_retriever_registry,
|
||||
reciprocal_rank_fusion,
|
||||
merge_adjacent_chunks,
|
||||
)
|
||||
|
||||
|
||||
class TestRetrievalResult:
|
||||
"""Tests for RetrievalResult dataclass."""
|
||||
|
||||
def test_result_creation(self):
|
||||
"""Test basic result creation."""
|
||||
result = RetrievalResult(text="Hello world", score=0.95)
|
||||
assert result.text == "Hello world"
|
||||
assert result.score == 0.95
|
||||
assert result.metadata == {}
|
||||
assert result.doc_id is None
|
||||
|
||||
def test_result_with_metadata(self):
|
||||
"""Test result with metadata."""
|
||||
result = RetrievalResult(
|
||||
text="Test content",
|
||||
score=0.8,
|
||||
metadata={"source": "test.txt"},
|
||||
doc_id="doc-1",
|
||||
chunk_index=2
|
||||
)
|
||||
assert result.metadata["source"] == "test.txt"
|
||||
assert result.doc_id == "doc-1"
|
||||
assert result.chunk_index == 2
|
||||
|
||||
def test_result_to_dict(self):
|
||||
"""Test result serialization."""
|
||||
result = RetrievalResult(
|
||||
text="Test",
|
||||
score=0.9,
|
||||
doc_id="doc-1"
|
||||
)
|
||||
d = result.to_dict()
|
||||
assert d["text"] == "Test"
|
||||
assert d["score"] == 0.9
|
||||
assert d["doc_id"] == "doc-1"
|
||||
|
||||
|
||||
class TestRetrievalStrategy:
|
||||
"""Tests for RetrievalStrategy enum."""
|
||||
|
||||
def test_strategy_values(self):
|
||||
"""Test strategy enum values."""
|
||||
assert RetrievalStrategy.BASIC == "basic"
|
||||
assert RetrievalStrategy.FUSION == "fusion"
|
||||
assert RetrievalStrategy.RECURSIVE == "recursive"
|
||||
assert RetrievalStrategy.AUTO_MERGE == "auto_merge"
|
||||
assert RetrievalStrategy.HYBRID == "hybrid"
|
||||
|
||||
def test_strategy_from_string(self):
|
||||
"""Test creating strategy from string."""
|
||||
assert RetrievalStrategy("basic") == RetrievalStrategy.BASIC
|
||||
assert RetrievalStrategy("fusion") == RetrievalStrategy.FUSION
|
||||
|
||||
|
||||
class TestReciprocalRankFusion:
|
||||
"""Tests for RRF algorithm."""
|
||||
|
||||
def test_single_list(self):
|
||||
"""Test RRF with single result list."""
|
||||
results = [
|
||||
RetrievalResult(text="A", score=0.9, doc_id="1"),
|
||||
RetrievalResult(text="B", score=0.8, doc_id="2"),
|
||||
]
|
||||
|
||||
fused = reciprocal_rank_fusion([results])
|
||||
assert len(fused) == 2
|
||||
assert fused[0].doc_id == "1" # First in original list
|
||||
|
||||
def test_multiple_lists(self):
|
||||
"""Test RRF with multiple result lists."""
|
||||
list1 = [
|
||||
RetrievalResult(text="A", score=0.9, doc_id="1"),
|
||||
RetrievalResult(text="B", score=0.8, doc_id="2"),
|
||||
]
|
||||
list2 = [
|
||||
RetrievalResult(text="B", score=0.95, doc_id="2"),
|
||||
RetrievalResult(text="C", score=0.7, doc_id="3"),
|
||||
]
|
||||
|
||||
fused = reciprocal_rank_fusion([list1, list2])
|
||||
|
||||
# B appears in both lists, should rank higher
|
||||
assert len(fused) == 3
|
||||
doc_ids = [r.doc_id for r in fused]
|
||||
assert "2" in doc_ids # B should be present
|
||||
|
||||
def test_empty_lists(self):
|
||||
"""Test RRF with empty lists."""
|
||||
fused = reciprocal_rank_fusion([])
|
||||
assert fused == []
|
||||
|
||||
fused = reciprocal_rank_fusion([[]])
|
||||
assert fused == []
|
||||
|
||||
def test_rrf_scores(self):
|
||||
"""Test that RRF scores are calculated correctly."""
|
||||
list1 = [
|
||||
RetrievalResult(text="A", score=0.9, doc_id="1"),
|
||||
]
|
||||
list2 = [
|
||||
RetrievalResult(text="A", score=0.8, doc_id="1"),
|
||||
]
|
||||
|
||||
fused = reciprocal_rank_fusion([list1, list2], k=60)
|
||||
|
||||
# A appears first in both lists, RRF score = 2 * (1 / (60 + 1))
|
||||
expected_score = 2 * (1 / 61)
|
||||
assert len(fused) == 1
|
||||
assert abs(fused[0].score - expected_score) < 0.001
|
||||
|
||||
|
||||
class TestMergeAdjacentChunks:
|
||||
"""Tests for chunk merging."""
|
||||
|
||||
def test_no_merge_needed(self):
|
||||
"""Test when no merging is needed."""
|
||||
results = [
|
||||
RetrievalResult(text="A", score=0.9, doc_id="1", chunk_index=0),
|
||||
RetrievalResult(text="B", score=0.8, doc_id="2", chunk_index=0),
|
||||
]
|
||||
|
||||
merged = merge_adjacent_chunks(results)
|
||||
assert len(merged) == 2
|
||||
|
||||
def test_merge_adjacent(self):
|
||||
"""Test merging adjacent chunks from same document."""
|
||||
results = [
|
||||
RetrievalResult(text="Part 1", score=0.9, doc_id="1", chunk_index=0),
|
||||
RetrievalResult(text="Part 2", score=0.8, doc_id="1", chunk_index=1),
|
||||
RetrievalResult(text="Other", score=0.7, doc_id="2", chunk_index=0),
|
||||
]
|
||||
|
||||
merged = merge_adjacent_chunks(results, max_gap=1)
|
||||
|
||||
# Should merge chunks 0 and 1 from doc 1
|
||||
assert len(merged) == 2
|
||||
|
||||
# Find the merged result
|
||||
merged_doc1 = [r for r in merged if r.doc_id == "1"][0]
|
||||
assert "Part 1" in merged_doc1.text
|
||||
assert "Part 2" in merged_doc1.text
|
||||
|
||||
def test_no_merge_with_gap(self):
|
||||
"""Test no merging when gap is too large."""
|
||||
results = [
|
||||
RetrievalResult(text="Part 1", score=0.9, doc_id="1", chunk_index=0),
|
||||
RetrievalResult(text="Part 5", score=0.8, doc_id="1", chunk_index=5),
|
||||
]
|
||||
|
||||
merged = merge_adjacent_chunks(results, max_gap=1)
|
||||
assert len(merged) == 2 # Not merged due to gap
|
||||
|
||||
def test_empty_results(self):
|
||||
"""Test with empty results."""
|
||||
merged = merge_adjacent_chunks([])
|
||||
assert merged == []
|
||||
|
||||
def test_no_chunk_index(self):
|
||||
"""Test results without chunk index."""
|
||||
results = [
|
||||
RetrievalResult(text="A", score=0.9, doc_id="1"),
|
||||
RetrievalResult(text="B", score=0.8, doc_id="1"),
|
||||
]
|
||||
|
||||
merged = merge_adjacent_chunks(results)
|
||||
assert len(merged) == 2 # Can't merge without chunk index
|
||||
|
||||
|
||||
class TestRetrieverRegistry:
|
||||
"""Tests for RetrieverRegistry."""
|
||||
|
||||
def test_list_retrievers(self):
|
||||
"""Test listing retrievers."""
|
||||
registry = get_retriever_registry()
|
||||
retrievers = registry.list_retrievers()
|
||||
assert isinstance(retrievers, list)
|
||||
|
||||
def test_get_nonexistent_retriever(self):
|
||||
"""Test getting non-existent retriever."""
|
||||
registry = get_retriever_registry()
|
||||
assert registry.get("nonexistent") is None
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Test that registry is a singleton."""
|
||||
registry1 = get_retriever_registry()
|
||||
registry2 = get_retriever_registry()
|
||||
assert registry1 is registry2
|
||||
@@ -0,0 +1,240 @@
|
||||
"""
|
||||
Unit tests for Vector Store Protocol and Registry.
|
||||
"""
|
||||
|
||||
from praisonaiagents.knowledge.vector_store import (
|
||||
VectorRecord,
|
||||
get_vector_store_registry,
|
||||
InMemoryVectorStore,
|
||||
)
|
||||
|
||||
|
||||
class TestVectorRecord:
|
||||
"""Tests for VectorRecord dataclass."""
|
||||
|
||||
def test_record_creation(self):
|
||||
"""Test basic record creation."""
|
||||
record = VectorRecord(
|
||||
id="test-id",
|
||||
text="Hello world",
|
||||
embedding=[0.1, 0.2, 0.3]
|
||||
)
|
||||
assert record.id == "test-id"
|
||||
assert record.text == "Hello world"
|
||||
assert record.embedding == [0.1, 0.2, 0.3]
|
||||
assert record.metadata == {}
|
||||
assert record.score is None
|
||||
|
||||
def test_record_with_metadata(self):
|
||||
"""Test record with metadata."""
|
||||
record = VectorRecord(
|
||||
id="test-id",
|
||||
text="Test content",
|
||||
embedding=[0.1],
|
||||
metadata={"source": "test.txt"}
|
||||
)
|
||||
assert record.metadata["source"] == "test.txt"
|
||||
|
||||
def test_record_to_dict(self):
|
||||
"""Test record serialization."""
|
||||
record = VectorRecord(
|
||||
id="test-id",
|
||||
text="Test",
|
||||
embedding=[0.1, 0.2],
|
||||
score=0.95
|
||||
)
|
||||
d = record.to_dict()
|
||||
assert d["id"] == "test-id"
|
||||
assert d["text"] == "Test"
|
||||
assert d["score"] == 0.95
|
||||
|
||||
def test_record_from_dict(self):
|
||||
"""Test record deserialization."""
|
||||
data = {
|
||||
"id": "test-id",
|
||||
"text": "Test content",
|
||||
"embedding": [0.1, 0.2, 0.3],
|
||||
"metadata": {"key": "value"},
|
||||
"score": 0.9
|
||||
}
|
||||
record = VectorRecord.from_dict(data)
|
||||
assert record.id == "test-id"
|
||||
assert record.score == 0.9
|
||||
|
||||
|
||||
class TestInMemoryVectorStore:
|
||||
"""Tests for InMemoryVectorStore."""
|
||||
|
||||
def test_add_vectors(self):
|
||||
"""Test adding vectors."""
|
||||
store = InMemoryVectorStore()
|
||||
ids = store.add(
|
||||
texts=["Hello", "World"],
|
||||
embeddings=[[0.1, 0.2], [0.3, 0.4]]
|
||||
)
|
||||
assert len(ids) == 2
|
||||
assert store.count() == 2
|
||||
|
||||
def test_add_with_metadata(self):
|
||||
"""Test adding vectors with metadata."""
|
||||
store = InMemoryVectorStore()
|
||||
ids = store.add(
|
||||
texts=["Test"],
|
||||
embeddings=[[0.1, 0.2]],
|
||||
metadatas=[{"source": "test.txt"}]
|
||||
)
|
||||
|
||||
records = store.get(ids)
|
||||
assert len(records) == 1
|
||||
assert records[0].metadata["source"] == "test.txt"
|
||||
|
||||
def test_query_similarity(self):
|
||||
"""Test querying by similarity."""
|
||||
store = InMemoryVectorStore()
|
||||
store.add(
|
||||
texts=["Hello world", "Goodbye world", "Hello there"],
|
||||
embeddings=[
|
||||
[1.0, 0.0, 0.0], # Most similar to query
|
||||
[0.0, 1.0, 0.0],
|
||||
[0.9, 0.1, 0.0] # Second most similar
|
||||
]
|
||||
)
|
||||
|
||||
# Query with embedding similar to first document
|
||||
results = store.query(
|
||||
embedding=[1.0, 0.0, 0.0],
|
||||
top_k=2
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0].text == "Hello world"
|
||||
assert results[0].score > results[1].score
|
||||
|
||||
def test_query_with_filter(self):
|
||||
"""Test querying with metadata filter."""
|
||||
store = InMemoryVectorStore()
|
||||
store.add(
|
||||
texts=["Doc A", "Doc B", "Doc C"],
|
||||
embeddings=[[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
|
||||
metadatas=[
|
||||
{"type": "a"},
|
||||
{"type": "b"},
|
||||
{"type": "a"}
|
||||
]
|
||||
)
|
||||
|
||||
results = store.query(
|
||||
embedding=[0.1, 0.2],
|
||||
top_k=10,
|
||||
filter={"type": "a"}
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
for r in results:
|
||||
assert r.metadata["type"] == "a"
|
||||
|
||||
def test_delete_by_ids(self):
|
||||
"""Test deleting by IDs."""
|
||||
store = InMemoryVectorStore()
|
||||
ids = store.add(
|
||||
texts=["A", "B", "C"],
|
||||
embeddings=[[0.1], [0.2], [0.3]]
|
||||
)
|
||||
|
||||
deleted = store.delete(ids=[ids[0], ids[1]])
|
||||
assert deleted == 2
|
||||
assert store.count() == 1
|
||||
|
||||
def test_delete_all(self):
|
||||
"""Test deleting all vectors."""
|
||||
store = InMemoryVectorStore()
|
||||
store.add(
|
||||
texts=["A", "B", "C"],
|
||||
embeddings=[[0.1], [0.2], [0.3]]
|
||||
)
|
||||
|
||||
deleted = store.delete(delete_all=True)
|
||||
assert deleted == 3
|
||||
assert store.count() == 0
|
||||
|
||||
def test_delete_by_filter(self):
|
||||
"""Test deleting by filter."""
|
||||
store = InMemoryVectorStore()
|
||||
store.add(
|
||||
texts=["A", "B", "C"],
|
||||
embeddings=[[0.1], [0.2], [0.3]],
|
||||
metadatas=[{"keep": True}, {"keep": False}, {"keep": True}]
|
||||
)
|
||||
|
||||
deleted = store.delete(filter={"keep": False})
|
||||
assert deleted == 1
|
||||
assert store.count() == 2
|
||||
|
||||
def test_get_by_ids(self):
|
||||
"""Test getting vectors by ID."""
|
||||
store = InMemoryVectorStore()
|
||||
ids = store.add(
|
||||
texts=["A", "B", "C"],
|
||||
embeddings=[[0.1], [0.2], [0.3]]
|
||||
)
|
||||
|
||||
records = store.get([ids[0], ids[2]])
|
||||
assert len(records) == 2
|
||||
texts = [r.text for r in records]
|
||||
assert "A" in texts
|
||||
assert "C" in texts
|
||||
|
||||
def test_namespace_isolation(self):
|
||||
"""Test namespace isolation."""
|
||||
store = InMemoryVectorStore()
|
||||
|
||||
store.add(texts=["NS1"], embeddings=[[0.1]], namespace="ns1")
|
||||
store.add(texts=["NS2"], embeddings=[[0.2]], namespace="ns2")
|
||||
|
||||
assert store.count(namespace="ns1") == 1
|
||||
assert store.count(namespace="ns2") == 1
|
||||
|
||||
results = store.query(embedding=[0.1], namespace="ns1")
|
||||
assert len(results) == 1
|
||||
assert results[0].text == "NS1"
|
||||
|
||||
def test_cosine_similarity(self):
|
||||
"""Test cosine similarity calculation."""
|
||||
# Identical vectors should have similarity 1.0
|
||||
sim = InMemoryVectorStore._cosine_similarity([1, 0, 0], [1, 0, 0])
|
||||
assert abs(sim - 1.0) < 0.001
|
||||
|
||||
# Orthogonal vectors should have similarity 0.0
|
||||
sim = InMemoryVectorStore._cosine_similarity([1, 0, 0], [0, 1, 0])
|
||||
assert abs(sim - 0.0) < 0.001
|
||||
|
||||
# Opposite vectors should have similarity -1.0
|
||||
sim = InMemoryVectorStore._cosine_similarity([1, 0, 0], [-1, 0, 0])
|
||||
assert abs(sim - (-1.0)) < 0.001
|
||||
|
||||
|
||||
class TestVectorStoreRegistry:
|
||||
"""Tests for VectorStoreRegistry."""
|
||||
|
||||
def test_memory_store_registered(self):
|
||||
"""Test that memory store is registered by default."""
|
||||
registry = get_vector_store_registry()
|
||||
assert "memory" in registry.list_stores()
|
||||
|
||||
def test_get_memory_store(self):
|
||||
"""Test getting memory store."""
|
||||
registry = get_vector_store_registry()
|
||||
store = registry.get("memory")
|
||||
assert store is not None
|
||||
assert store.name == "memory"
|
||||
|
||||
def test_get_nonexistent_store(self):
|
||||
"""Test getting non-existent store."""
|
||||
registry = get_vector_store_registry()
|
||||
assert registry.get("nonexistent") is None
|
||||
|
||||
def test_singleton_pattern(self):
|
||||
"""Test that registry is a singleton."""
|
||||
registry1 = get_vector_store_registry()
|
||||
registry2 = get_vector_store_registry()
|
||||
assert registry1 is registry2
|
||||
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.3.0" },
|
||||
{ name = "praisonaiagents", specifier = ">=0.3.1" },
|
||||
{ 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.3.0"
|
||||
version = "0.3.1"
|
||||
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/fc/79/08c9f2d8c12f9521f9fcc3df2510bd0d09377fd68b27413c64e501c1d47b/praisonaiagents-0.3.0.tar.gz", hash = "sha256:6fac4b754ec924ba804ef09ebef9919f276861d9f6d8c66cc33721063cb0d04c", size = 741299, upload-time = "2025-12-26T02:18:18.235Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/dc/bd/57319c557108a64ada4294ca69a75a419050ee37dfc644c776435391ecf4/praisonaiagents-0.3.1.tar.gz", hash = "sha256:fb023935176c1d16f110b4356f7327fd04092d01448845ad4e0fc66882feeeff", size = 757190, upload-time = "2025-12-26T04:03:55.969Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/18/4e/2165feb9897d9b17522d7653b9111e40262bcf5172fc6fb549e5f0b2413f/praisonaiagents-0.3.0-py3-none-any.whl", hash = "sha256:fdbe3cf290e274ffd518b932d0ad8fd516b1600324bde2bf1970c5efd26c4292", size = 689391, upload-time = "2025-12-26T02:18:16.247Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/85/0c/0f11bbe77bd31b4870b5caa8a9c69b570a8556061afe1472b7350c98b66c/praisonaiagents-0.3.1-py3-none-any.whl", hash = "sha256:253f90a036cfb247d6fd532df7f2b8a022634dfeb83b1ce58c99d32ca7deb6e1", size = 707840, upload-time = "2025-12-26T04:03:54.013Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in new issue
Block a user