Compare commits
25 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
6158046eae
|
|||
|
ce948da5ff
|
|||
|
fea2cc2649
|
|||
|
7d646d3439
|
|||
|
14d390a660
|
|||
|
07f741fafa
|
|||
|
07d0bb9da6
|
|||
|
59ea0d2092
|
|||
|
a2c859aef6
|
|||
|
ac1e17b8bd
|
|||
|
527fccdf6b
|
|||
|
9e89c031fd
|
|||
|
82078bd9e8
|
|||
|
2f4f4c9b49
|
|||
|
ada913179d
|
|||
|
ca33ce2f70
|
|||
|
2c13e747b4
|
|||
|
a4918ef0e6
|
|||
|
e1e85ed9c7
|
|||
|
a1b65057a9
|
|||
|
d468e36e59
|
|||
|
a6e1151209
|
|||
|
444cbbfcc6
|
|||
|
7cb5a542cb
|
|||
|
cdf13f842d
|
@@ -90,7 +90,10 @@ celerybeat.pid
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.env.*
|
||||
!.env.example
|
||||
.venv
|
||||
.venv/
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
@@ -126,6 +129,7 @@ cython_debug/
|
||||
|
||||
# macOS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# uv
|
||||
uv.lock
|
||||
@@ -167,3 +171,8 @@ PIPE
|
||||
|
||||
# Git worktrees for parallel task branches
|
||||
worktrees/
|
||||
logs/
|
||||
data/
|
||||
|
||||
# Internal modification report (not for upstream)
|
||||
docs/camodif.md
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
name: local/fastapi-health
|
||||
description: Build a simple REST API with a /health endpoint in Python using FastAPI
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A main.py file exists with a FastAPI app that has a /health endpoint
|
||||
returning status ok. Requirements are listed in requirements.txt.
|
||||
@@ -0,0 +1,15 @@
|
||||
name: local/basic-rag
|
||||
description: >
|
||||
Build a basic RAG (Retrieval Augmented Generation) application in Python.
|
||||
Use OpenAI embeddings for vectorizing documents and FAISS for the vector store.
|
||||
Include a document loader that reads .txt files from a directory, chunks them,
|
||||
embeds them, stores in FAISS, and answers questions by retrieving relevant
|
||||
chunks and passing them to OpenAI GPT for answer generation. Include a CLI
|
||||
interface that accepts a question as input. Include requirements.txt.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A rag.py implements document loading, chunking, embedding with OpenAI,
|
||||
FAISS vector storage, retrieval, and answer generation.
|
||||
CLI accepts a question via argparse and a --docs-dir flag for the documents directory.
|
||||
A requirements.txt lists all dependencies (openai, faiss-cpu, tiktoken, etc).
|
||||
@@ -0,0 +1,11 @@
|
||||
name: local/todo-cli
|
||||
description: >
|
||||
Build a command-line TODO application in Python with add, list, complete,
|
||||
and delete commands. Store tasks in a JSON file. Include a Task dataclass
|
||||
with id, title, completed, and created_at fields.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A todo.py file exists with CLI commands for add, list, complete, and delete.
|
||||
Tasks are persisted in a tasks.json file. A Task dataclass has id, title,
|
||||
completed, and created_at fields.
|
||||
@@ -0,0 +1,14 @@
|
||||
name: local/csv-analyzer
|
||||
description: >
|
||||
Build a CSV data analysis tool in Python. Load CSV files, detect column types
|
||||
automatically, compute summary statistics (mean, median, mode, std dev, quartiles),
|
||||
detect outliers using IQR method, generate correlation matrices, and output
|
||||
a formatted text report. Support filtering rows by column conditions.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A loader.py handles CSV reading with automatic type detection and encoding handling.
|
||||
An analyzer.py computes statistics, outlier detection, and correlations.
|
||||
A reporter.py formats results into a readable text report.
|
||||
A cli.py provides argparse interface with --file, --columns, --filter, and --output flags.
|
||||
A requirements.txt lists dependencies.
|
||||
@@ -0,0 +1,16 @@
|
||||
name: local/rate-limiter
|
||||
description: >
|
||||
Build a rate limiting middleware library in Python for Flask APIs.
|
||||
Implement token bucket, sliding window, and fixed window algorithms.
|
||||
Support per-IP and per-API-key limits. Store counters in an in-memory
|
||||
store with TTL expiry. Include a Flask decorator for easy route-level
|
||||
rate limiting and configurable response headers (X-RateLimit-*).
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
An algorithms.py implements TokenBucket, SlidingWindow, and FixedWindow classes.
|
||||
A store.py provides an in-memory counter store with TTL-based expiry.
|
||||
A middleware.py provides a Flask decorator @rate_limit() for route protection.
|
||||
A config.py defines rate limit configuration dataclasses.
|
||||
An example_app.py shows a Flask app using all three algorithms on different routes.
|
||||
A requirements.txt lists dependencies (flask).
|
||||
@@ -0,0 +1,15 @@
|
||||
name: local/json-validator
|
||||
description: >
|
||||
Build a JSON schema validator in Python from scratch (no jsonschema library).
|
||||
Support type validation (string, number, integer, boolean, array, object, null),
|
||||
constraints (minLength, maxLength, minimum, maximum, pattern, enum),
|
||||
nested object and array validation, required fields, and custom error messages.
|
||||
Output validation errors with JSON path to the failing field.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A validator.py implements SchemaValidator with recursive type and constraint checking.
|
||||
A types.py defines validation rules for each JSON type with constraint support.
|
||||
An errors.py provides ValidationError with JSON path tracking and human-readable messages.
|
||||
A cli.py accepts --schema and --data file paths via argparse and prints validation results.
|
||||
A requirements.txt lists dependencies (none beyond stdlib expected).
|
||||
@@ -0,0 +1,16 @@
|
||||
name: local/api-tester
|
||||
description: >
|
||||
Build an HTTP API testing framework in Python. Define test cases in YAML files
|
||||
specifying method, URL, headers, body, and expected response (status code,
|
||||
body contains, JSON path assertions). Run tests sequentially, support variable
|
||||
extraction from responses for chaining requests, and output results with
|
||||
pass/fail counts and detailed failure messages.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A runner.py loads YAML test suites and executes HTTP requests sequentially.
|
||||
An assertions.py validates responses against expected status, body, and JSON path rules.
|
||||
A variables.py handles variable extraction from responses and substitution in subsequent requests.
|
||||
A reporter.py outputs test results with pass/fail summary and failure details.
|
||||
A cli.py provides argparse interface with --suite (YAML path) and --base-url flags.
|
||||
A requirements.txt lists dependencies (requests, pyyaml).
|
||||
@@ -0,0 +1,13 @@
|
||||
name: local/bookstore-api
|
||||
description: >
|
||||
Build a REST API for a bookstore using FastAPI. Include endpoints for
|
||||
CRUD operations on books (GET /books, GET /books/{id}, POST /books,
|
||||
PUT /books/{id}, DELETE /books/{id}). Use Pydantic models for Book
|
||||
with fields: id, title, author, price, isbn. Store books in memory.
|
||||
Include requirements.txt.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A main.py with FastAPI app has all CRUD endpoints for books.
|
||||
A Book Pydantic model has id, title, author, price, isbn fields.
|
||||
A requirements.txt lists dependencies.
|
||||
@@ -0,0 +1,12 @@
|
||||
name: local/websocket-chat
|
||||
description: >
|
||||
Build a WebSocket chat server in Python using the websockets library.
|
||||
Support multiple connected clients, broadcast messages to all clients,
|
||||
handle connect/disconnect events, and include a simple HTML client page.
|
||||
Include requirements.txt.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A server.py implements a WebSocket chat server with broadcast support.
|
||||
An index.html provides a browser-based chat client.
|
||||
A requirements.txt lists dependencies.
|
||||
@@ -0,0 +1,12 @@
|
||||
name: local/hn-scraper
|
||||
description: >
|
||||
Build a Hacker News scraper in Python using requests and BeautifulSoup.
|
||||
Scrape the front page for top stories including title, URL, points, and
|
||||
comment count. Output results as JSON. Include a Story dataclass and
|
||||
requirements.txt.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A scraper.py fetches and parses Hacker News front page stories.
|
||||
A Story dataclass has title, url, points, and comments fields.
|
||||
Results are output as JSON. A requirements.txt lists dependencies.
|
||||
@@ -0,0 +1,12 @@
|
||||
name: local/flask-auth
|
||||
description: >
|
||||
Build a Flask web application with user authentication. Include register
|
||||
and login endpoints, password hashing with bcrypt, JWT token generation,
|
||||
and a protected endpoint that requires authentication. Use SQLite for
|
||||
user storage. Include requirements.txt.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
An app.py has Flask routes for register, login, and a protected endpoint.
|
||||
Passwords are hashed with bcrypt. JWT tokens are used for auth.
|
||||
SQLite stores users. A requirements.txt lists dependencies.
|
||||
@@ -0,0 +1,15 @@
|
||||
name: local/weather-cli
|
||||
description: >
|
||||
Build a Python CLI weather dashboard that fetches weather data from
|
||||
the OpenWeatherMap API. Include a WeatherData dataclass, formatted
|
||||
terminal output with current conditions (temperature, humidity, wind),
|
||||
argument parsing for city name, and proper error handling for API
|
||||
failures. Include requirements.txt.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
A weather.py fetches weather data from OpenWeatherMap API for a given city.
|
||||
A WeatherData dataclass holds temperature, humidity, wind_speed, and description.
|
||||
CLI accepts city name via argparse. Output is formatted for terminal display.
|
||||
Error handling covers network failures and invalid cities.
|
||||
A requirements.txt lists dependencies.
|
||||
@@ -0,0 +1,15 @@
|
||||
name: local/url-shortener
|
||||
description: >
|
||||
Build a URL shortener service in Python using Flask. Include endpoints
|
||||
to shorten a URL (POST /shorten) and redirect via short code (GET /<code>).
|
||||
Use an in-memory dictionary for storage. Generate short codes using
|
||||
base62 encoding of an auto-incrementing counter. Include input validation
|
||||
for URLs, a simple HTML landing page, and requirements.txt.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
An app.py has Flask routes for POST /shorten and GET /<code>.
|
||||
Short codes are generated via base62 encoding.
|
||||
URL validation rejects malformed inputs.
|
||||
An index.html landing page allows users to paste and shorten URLs.
|
||||
In-memory dict stores mappings. A requirements.txt lists dependencies.
|
||||
@@ -0,0 +1,21 @@
|
||||
name: local/rag-multifile
|
||||
description: >
|
||||
Build a production-quality RAG (Retrieval Augmented Generation) application in Python.
|
||||
Use OpenAI embeddings for vectorizing documents and FAISS for the vector store.
|
||||
Follow a clean project structure with separate modules for each concern.
|
||||
The project should have a src/ subdirectory containing all modules, a config
|
||||
file for settings, and a top-level main entry point.
|
||||
strategy_actor: anthropic/claude-sonnet-4-20250514
|
||||
execution_actor: anthropic/claude-sonnet-4-20250514
|
||||
definition_of_done: >
|
||||
Project structure must follow this layout:
|
||||
src/document_loader.py — reads .txt files from a directory and returns raw text.
|
||||
src/chunker.py — splits documents into overlapping chunks using tiktoken tokenization.
|
||||
src/embeddings.py — wraps OpenAI embedding API calls with batching support.
|
||||
src/vector_store.py — manages FAISS index creation, add, search, and save/load to disk.
|
||||
src/retriever.py — combines embedding + vector search to find relevant chunks for a query.
|
||||
src/generator.py — sends retrieved context + question to OpenAI GPT and returns the answer.
|
||||
src/__init__.py — empty init file for the src package.
|
||||
config.py — dataclass-based configuration (model names, chunk size, overlap, top-k).
|
||||
main.py — CLI entry point with argparse: --docs-dir, --question, --interactive, --model flags.
|
||||
requirements.txt — lists all dependencies (openai, faiss-cpu, tiktoken, numpy, python-dotenv).
|
||||
+161
@@ -0,0 +1,161 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run all 7 simulations through the V3 lifecycle, collecting files after each."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
CLI = [sys.executable, "-m", "cleveragents.cli.main"]
|
||||
PROJECT_ROOT = Path.cwd()
|
||||
|
||||
SIMS = [
|
||||
("sim1", "local/todo-cli"),
|
||||
("sim2", "local/bookstore-api"),
|
||||
("sim3", "local/websocket-chat"),
|
||||
("sim4", "local/hn-scraper"),
|
||||
("sim5", "local/flask-auth"),
|
||||
("sim6", "local/weather-cli"),
|
||||
("sim7", "local/url-shortener"),
|
||||
]
|
||||
|
||||
# Track files that existed before we started
|
||||
IGNORE_PATTERNS = {
|
||||
".venv", ".git", "src", "docs", "actions", "simulations", ".cleveragents",
|
||||
".claude", ".env", ".gitignore", "run_all_sims.py", "features", "tests",
|
||||
"benchmarks", "pyproject.toml", "poetry.lock", "alembic.ini", "Makefile",
|
||||
"Dockerfile", "docker-compose.yml", "README.md", "CLAUDE.md", "LICENSE",
|
||||
"cleveragents.db", "__pycache__",
|
||||
}
|
||||
|
||||
|
||||
def run(args, timeout=300):
|
||||
r = subprocess.run(CLI + args, capture_output=True, text=True, timeout=timeout)
|
||||
return r.stdout + r.stderr
|
||||
|
||||
|
||||
def extract_plan_id(output):
|
||||
for line in output.splitlines():
|
||||
if "ID:" in line:
|
||||
parts = line.split("ID:")
|
||||
if len(parts) > 1:
|
||||
return parts[1].strip().strip("│").strip()
|
||||
return None
|
||||
|
||||
|
||||
def extract_file_count(output):
|
||||
for line in output.splitlines():
|
||||
if "file_count=" in line:
|
||||
for part in line.split():
|
||||
if part.startswith("file_count="):
|
||||
return int(part.split("=")[1])
|
||||
return 0
|
||||
|
||||
|
||||
def snapshot_cwd():
|
||||
"""Return set of all files in CWD (relative paths), excluding known dirs."""
|
||||
files = set()
|
||||
for item in PROJECT_ROOT.iterdir():
|
||||
if item.name in IGNORE_PATTERNS:
|
||||
continue
|
||||
if item.is_file():
|
||||
files.add(item.name)
|
||||
elif item.is_dir():
|
||||
for f in item.rglob("*"):
|
||||
if f.is_file():
|
||||
files.add(str(f.relative_to(PROJECT_ROOT)))
|
||||
return files
|
||||
|
||||
|
||||
def collect_new_files(before_snapshot, sim_dir):
|
||||
"""Move newly created files to sim_dir."""
|
||||
after = snapshot_cwd()
|
||||
new_files = after - before_snapshot
|
||||
sim_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for rel_path in sorted(new_files):
|
||||
src = PROJECT_ROOT / rel_path
|
||||
dst = sim_dir / rel_path
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.move(str(src), str(dst))
|
||||
|
||||
# Clean up any empty dirs left behind
|
||||
for rel_path in sorted(new_files, reverse=True):
|
||||
src_parent = (PROJECT_ROOT / rel_path).parent
|
||||
if src_parent != PROJECT_ROOT and src_parent.exists():
|
||||
try:
|
||||
src_parent.rmdir() # only removes if empty
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
return sorted(new_files)
|
||||
|
||||
|
||||
results = {}
|
||||
|
||||
for sim_name, action_name in SIMS:
|
||||
print(f"\n{'='*60}")
|
||||
print(f" {sim_name.upper()}: {action_name}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
sim_dir = PROJECT_ROOT / "simulations" / sim_name
|
||||
if sim_dir.exists():
|
||||
shutil.rmtree(sim_dir)
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
# Plan use
|
||||
out = run(["plan", "use", action_name])
|
||||
plan_id = extract_plan_id(out)
|
||||
if not plan_id:
|
||||
print(f" FAILED: Could not extract plan ID")
|
||||
results[sim_name] = {"status": "FAILED", "error": "No plan ID"}
|
||||
continue
|
||||
print(f" Plan ID: {plan_id}")
|
||||
|
||||
# Execute
|
||||
print(f" Executing...", end="", flush=True)
|
||||
before = snapshot_cwd()
|
||||
out = run(["plan", "execute", plan_id], timeout=300)
|
||||
elapsed = time.time() - t0
|
||||
file_count = extract_file_count(out)
|
||||
|
||||
if "Execute completed" not in out and "Plan Executed" not in out:
|
||||
print(f" FAILED ({elapsed:.0f}s)")
|
||||
results[sim_name] = {"status": "FAILED", "error": out[-300:]}
|
||||
continue
|
||||
print(f" done ({elapsed:.0f}s, {file_count} files in sandbox)")
|
||||
|
||||
# Apply
|
||||
out = run(["plan", "lifecycle-apply", plan_id])
|
||||
if "file(s) written" not in out:
|
||||
print(f" Apply FAILED")
|
||||
results[sim_name] = {"status": "FAILED", "error": "Apply failed"}
|
||||
continue
|
||||
|
||||
# Collect files
|
||||
new_files = collect_new_files(before, sim_dir)
|
||||
print(f" Applied {len(new_files)} files:")
|
||||
for f in new_files:
|
||||
print(f" {f}")
|
||||
|
||||
results[sim_name] = {
|
||||
"status": "PASS",
|
||||
"plan_id": plan_id,
|
||||
"sandbox_count": file_count,
|
||||
"applied_files": new_files,
|
||||
"elapsed": f"{elapsed:.0f}s",
|
||||
}
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(" SUMMARY")
|
||||
print(f"{'='*60}")
|
||||
for sim, info in results.items():
|
||||
status = info["status"]
|
||||
files = info.get("applied_files", [])
|
||||
elapsed = info.get("elapsed", "?")
|
||||
print(f" {sim}: {status} — {len(files)} files ({elapsed})")
|
||||
for f in files:
|
||||
print(f" {f}")
|
||||
@@ -0,0 +1,316 @@
|
||||
I'll create a production-ready CLI todo application using Click. Here's the complete implementation:
|
||||
|
||||
## `todo.py`
|
||||
```python
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
A simple CLI todo application built with Click.
|
||||
"""
|
||||
import json
|
||||
import click
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from typing import List, Dict, Any
|
||||
import sys
|
||||
|
||||
|
||||
class TodoManager:
|
||||
"""Manages todo items with JSON file persistence."""
|
||||
|
||||
def __init__(self, data_file: str = "todos.json"):
|
||||
self.data_file = Path(data_file)
|
||||
self._ensure_data_file()
|
||||
|
||||
def _ensure_data_file(self) -> None:
|
||||
"""Create data file if it doesn't exist."""
|
||||
if not self.data_file.exists():
|
||||
self._save_todos([])
|
||||
|
||||
def _load_todos(self) -> List[Dict[str, Any]]:
|
||||
"""Load todos from JSON file."""
|
||||
try:
|
||||
with open(self.data_file, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, FileNotFoundError) as e:
|
||||
click.echo(f"Error reading todos: {e}", err=True)
|
||||
return []
|
||||
|
||||
def _save_todos(self, todos: List[Dict[str, Any]]) -> None:
|
||||
"""Save todos to JSON file."""
|
||||
try:
|
||||
with open(self.data_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(todos, f, indent=2, ensure_ascii=False)
|
||||
except IOError as e:
|
||||
click.echo(f"Error saving todos: {e}", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
def add_todo(self, text: str) -> int:
|
||||
"""Add a new todo item and return its ID."""
|
||||
todos = self._load_todos()
|
||||
todo_id = max([todo.get('id', 0) for todo in todos], default=0) + 1
|
||||
|
||||
new_todo = {
|
||||
'id': todo_id,
|
||||
'text': text.strip(),
|
||||
'created': datetime.now().isoformat(),
|
||||
'completed': False
|
||||
}
|
||||
|
||||
todos.append(new_todo)
|
||||
self._save_todos(todos)
|
||||
return todo_id
|
||||
|
||||
def list_todos(self) -> List[Dict[str, Any]]:
|
||||
"""Get all todo items."""
|
||||
return self._load_todos()
|
||||
|
||||
def delete_todo(self, todo_id: int) -> bool:
|
||||
"""Delete a todo item by ID. Returns True if deleted, False if not found."""
|
||||
todos = self._load_todos()
|
||||
original_length = len(todos)
|
||||
todos = [todo for todo in todos if todo.get('id') != todo_id]
|
||||
|
||||
if len(todos) < original_length:
|
||||
self._save_todos(todos)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# Global todo manager instance
|
||||
todo_manager = TodoManager()
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.version_option(version="1.0.0")
|
||||
def cli():
|
||||
"""A simple CLI todo application.
|
||||
|
||||
Manage your todo items with add, list, and delete commands.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.argument('text', required=True)
|
||||
def add(text: str):
|
||||
"""Add a new todo item.
|
||||
|
||||
TEXT: The todo item description
|
||||
|
||||
Example: todo add "Buy groceries"
|
||||
"""
|
||||
if not text.strip():
|
||||
click.echo("Error: Todo text cannot be empty", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
todo_id = todo_manager.add_todo(text)
|
||||
click.echo(f"✅ Added todo #{todo_id}: {text}")
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.option('--show-ids', is_flag=True, help='Show todo IDs for deletion')
|
||||
def list(show_ids: bool):
|
||||
"""List all todo items.
|
||||
|
||||
Use --show-ids to display IDs for deletion.
|
||||
|
||||
Example: todo list --show-ids
|
||||
"""
|
||||
todos = todo_manager.list_todos()
|
||||
|
||||
if not todos:
|
||||
click.echo("📝 No todos found. Add some with 'todo add <text>'")
|
||||
return
|
||||
|
||||
click.echo(f"📋 You have {len(todos)} todo(s):\n")
|
||||
|
||||
for todo in todos:
|
||||
status = "✅" if todo.get('completed') else "⏳"
|
||||
text = todo.get('text', '')
|
||||
|
||||
if show_ids:
|
||||
click.echo(f" {status} #{todo.get('id')}: {text}")
|
||||
else:
|
||||
click.echo(f" {status} {text}")
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.argument('todo_id', type=int)
|
||||
@click.confirmation_option(prompt='Are you sure you want to delete this todo?')
|
||||
def delete(todo_id: int):
|
||||
"""Delete a todo item by ID.
|
||||
|
||||
TODO_ID: The ID of the todo item to delete (use 'todo list --show-ids' to see IDs)
|
||||
|
||||
Example: todo delete 1
|
||||
"""
|
||||
if todo_manager.delete_todo(todo_id):
|
||||
click.echo(f"🗑️ Deleted todo #{todo_id}")
|
||||
else:
|
||||
click.echo(f"❌ Todo #{todo_id} not found", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
cli()
|
||||
```
|
||||
|
||||
## `requirements.txt`
|
||||
```txt
|
||||
click>=8.0.0
|
||||
```
|
||||
|
||||
## `setup.py`
|
||||
```python
|
||||
from setuptools import setup, find_packages
|
||||
|
||||
setup(
|
||||
name='todo-cli',
|
||||
version='1.0.0',
|
||||
py_modules=['todo'],
|
||||
install_requires=[
|
||||
'click>=8.0.0',
|
||||
],
|
||||
entry_points={
|
||||
'console_scripts': [
|
||||
'todo=todo:cli',
|
||||
],
|
||||
},
|
||||
author='Your Name',
|
||||
author_email='your.email@example.com',
|
||||
description='A simple CLI todo application',
|
||||
long_description=open('README.md').read(),
|
||||
long_description_content_type='text/markdown',
|
||||
url='https://github.com/yourusername/todo-cli',
|
||||
classifiers=[
|
||||
'Development Status :: 4 - Beta',
|
||||
'Intended Audience :: Developers',
|
||||
'License :: OSI Approved :: MIT License',
|
||||
'Programming Language :: Python :: 3',
|
||||
'Programming Language :: Python :: 3.7+',
|
||||
],
|
||||
python_requires='>=3.7',
|
||||
)
|
||||
```
|
||||
|
||||
## `README.md`
|
||||
```markdown
|
||||
# Todo CLI
|
||||
|
||||
A simple command-line todo application built with Python and Click.
|
||||
|
||||
## Features
|
||||
|
||||
- ✅ Add new todo items
|
||||
- 📋 List all todos with optional IDs
|
||||
- 🗑️ Delete todos by ID
|
||||
- 💾 Persistent storage using JSON
|
||||
- 🎨 Colorful CLI interface with emojis
|
||||
|
||||
## Installation
|
||||
|
||||
### Option 1: Install as a package
|
||||
```bash
|
||||
pip install -e .
|
||||
todo --help
|
||||
```
|
||||
|
||||
### Option 2: Run directly
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
python todo.py --help
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Add a todo
|
||||
```bash
|
||||
todo add "Buy groceries"
|
||||
todo add "Finish project documentation"
|
||||
```
|
||||
|
||||
### List todos
|
||||
```bash
|
||||
# Simple list
|
||||
todo list
|
||||
|
||||
# List with IDs (needed for deletion)
|
||||
todo list --show-ids
|
||||
```
|
||||
|
||||
### Delete a todo
|
||||
```bash
|
||||
# First, get the ID with --show-ids
|
||||
todo list --show-ids
|
||||
|
||||
# Then delete by ID
|
||||
todo delete 1
|
||||
```
|
||||
|
||||
### Get help
|
||||
```bash
|
||||
todo --help
|
||||
todo add --help
|
||||
todo list --help
|
||||
todo delete --help
|
||||
```
|
||||
|
||||
## Data Storage
|
||||
|
||||
Todos are stored in a `todos.json` file in the current directory. The file is automatically created when you add your first todo.
|
||||
|
||||
## Example Session
|
||||
|
||||
```bash
|
||||
$ todo add "Learn Python"
|
||||
✅ Added todo #1: Learn Python
|
||||
|
||||
$ todo add "Build a CLI app"
|
||||
✅ Added todo #2: Build a CLI app
|
||||
|
||||
$ todo list --show-ids
|
||||
📋 You have 2 todo(s):
|
||||
|
||||
⏳ #1: Learn Python
|
||||
⏳ #2: Build a CLI app
|
||||
|
||||
$ todo delete 1
|
||||
Are you sure you want to delete this todo? [y/N]: y
|
||||
🗑️ Deleted todo #1
|
||||
|
||||
$ todo list
|
||||
📋 You have 1 todo(s):
|
||||
|
||||
⏳ Build a CLI app
|
||||
```
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.7+
|
||||
- Click 8.0+
|
||||
|
||||
## License
|
||||
|
||||
MIT License
|
||||
```
|
||||
|
||||
## Key Features
|
||||
|
||||
1. **Production-ready**: Error handling, input validation, and proper exit codes
|
||||
2. **User-friendly**: Confirmation prompts, helpful messages, and emoji indicators
|
||||
3. **Robust storage**: JSON persistence with proper file handling
|
||||
4. **Flexible installation**: Can be installed as a package or run directly
|
||||
5. **Good UX**: Clear help messages, optional ID display, and status indicators
|
||||
6. **Cross-platform**: Uses `pathlib` for proper file handling
|
||||
|
||||
## Usage Examples
|
||||
|
||||
```bash
|
||||
# Install and run
|
||||
pip install -r requirements.txt
|
||||
python todo.py add "My first todo"
|
||||
python todo.py list --show-ids
|
||||
python todo.py delete 1
|
||||
```
|
||||
|
||||
The application handles edge cases like empty todos, missing files, and invalid IDs gracefully with appropriate error messages.
|
||||
@@ -0,0 +1 @@
|
||||
[]
|
||||
@@ -0,0 +1,253 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass, asdict
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
id: str
|
||||
title: str
|
||||
completed: bool
|
||||
created_at: str
|
||||
|
||||
@classmethod
|
||||
def create(cls, title: str) -> 'Task':
|
||||
"""Create a new task with auto-generated ID and timestamp."""
|
||||
return cls(
|
||||
id=str(uuid.uuid4()),
|
||||
title=title,
|
||||
completed=False,
|
||||
created_at=datetime.now().isoformat()
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> 'Task':
|
||||
"""Create Task instance from dictionary."""
|
||||
return cls(**data)
|
||||
|
||||
|
||||
class TaskManager:
|
||||
def __init__(self, storage_file: str = "tasks.json"):
|
||||
self.storage_file = Path(storage_file)
|
||||
self.tasks: List[Task] = []
|
||||
self.load_tasks()
|
||||
|
||||
def load_tasks(self) -> None:
|
||||
"""Load tasks from JSON file."""
|
||||
try:
|
||||
if self.storage_file.exists():
|
||||
with open(self.storage_file, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
self.tasks = [Task.from_dict(task_dict) for task_dict in data]
|
||||
else:
|
||||
self.tasks = []
|
||||
except (json.JSONDecodeError, KeyError) as e:
|
||||
print(f"Error loading tasks: {e}")
|
||||
print("Starting with empty task list.")
|
||||
self.tasks = []
|
||||
|
||||
def save_tasks(self) -> None:
|
||||
"""Save tasks to JSON file."""
|
||||
try:
|
||||
with open(self.storage_file, 'w', encoding='utf-8') as f:
|
||||
json.dump([asdict(task) for task in self.tasks], f, indent=2, ensure_ascii=False)
|
||||
except IOError as e:
|
||||
print(f"Error saving tasks: {e}")
|
||||
|
||||
def add_task(self, title: str) -> Task:
|
||||
"""Add a new task."""
|
||||
task = Task.create(title)
|
||||
self.tasks.append(task)
|
||||
self.save_tasks()
|
||||
return task
|
||||
|
||||
def list_tasks(self) -> List[Task]:
|
||||
"""Get all tasks."""
|
||||
return self.tasks
|
||||
|
||||
def find_task(self, task_id: str) -> Optional[Task]:
|
||||
"""Find task by ID."""
|
||||
for task in self.tasks:
|
||||
if task.id == task_id:
|
||||
return task
|
||||
return None
|
||||
|
||||
def complete_task(self, task_id: str) -> bool:
|
||||
"""Mark task as completed."""
|
||||
task = self.find_task(task_id)
|
||||
if task:
|
||||
task.completed = True
|
||||
self.save_tasks()
|
||||
return True
|
||||
return False
|
||||
|
||||
def delete_task(self, task_id: str) -> bool:
|
||||
"""Delete task by ID."""
|
||||
task = self.find_task(task_id)
|
||||
if task:
|
||||
self.tasks.remove(task)
|
||||
self.save_tasks()
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def format_task(task: Task) -> str:
|
||||
"""Format task for display."""
|
||||
status = "✓" if task.completed else "○"
|
||||
created_date = datetime.fromisoformat(task.created_at).strftime("%Y-%m-%d %H:%M")
|
||||
return f"{status} [{task.id[:8]}] {task.title} (created: {created_date})"
|
||||
|
||||
|
||||
def cmd_add(args, task_manager: TaskManager) -> None:
|
||||
"""Handle add command."""
|
||||
if not args.title:
|
||||
print("Error: Task title cannot be empty.")
|
||||
return
|
||||
|
||||
task = task_manager.add_task(args.title)
|
||||
print(f"Added task: {task.title}")
|
||||
print(f"Task ID: {task.id}")
|
||||
|
||||
|
||||
def cmd_list(args, task_manager: TaskManager) -> None:
|
||||
"""Handle list command."""
|
||||
tasks = task_manager.list_tasks()
|
||||
|
||||
if not tasks:
|
||||
print("No tasks found.")
|
||||
return
|
||||
|
||||
print(f"Total tasks: {len(tasks)}")
|
||||
print()
|
||||
|
||||
for task in tasks:
|
||||
print(format_task(task))
|
||||
|
||||
|
||||
def cmd_complete(args, task_manager: TaskManager) -> None:
|
||||
"""Handle complete command."""
|
||||
if task_manager.complete_task(args.id):
|
||||
task = task_manager.find_task(args.id)
|
||||
print(f"Completed task: {task.title}")
|
||||
else:
|
||||
print(f"Error: Task with ID '{args.id}' not found.")
|
||||
|
||||
|
||||
def cmd_delete(args, task_manager: TaskManager) -> None:
|
||||
"""Handle delete command."""
|
||||
task = task_manager.find_task(args.id)
|
||||
if task and task_manager.delete_task(args.id):
|
||||
print(f"Deleted task: {task.title}")
|
||||
else:
|
||||
print(f"Error: Task with ID '{args.id}' not found.")
|
||||
|
||||
|
||||
def create_parser() -> argparse.ArgumentParser:
|
||||
"""Create and configure argument parser."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Simple command-line TODO application",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
%(prog)s add "Buy groceries"
|
||||
%(prog)s list
|
||||
%(prog)s complete abc12345
|
||||
%(prog)s delete abc12345
|
||||
"""
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest='command', help='Available commands')
|
||||
|
||||
# Add command
|
||||
add_parser = subparsers.add_parser('add', help='Add a new task')
|
||||
add_parser.add_argument('title', help='Task title')
|
||||
|
||||
# List command
|
||||
subparsers.add_parser('list', help='List all tasks')
|
||||
|
||||
# Complete command
|
||||
complete_parser = subparsers.add_parser('complete', help='Mark task as completed')
|
||||
complete_parser.add_argument('id', help='Task ID (or first 8 characters)')
|
||||
|
||||
# Delete command
|
||||
delete_parser = subparsers.add_parser('delete', help='Delete a task')
|
||||
delete_parser.add_argument('id', help='Task ID (or first 8 characters)')
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def find_task_by_partial_id(task_manager: TaskManager, partial_id: str) -> Optional[Task]:
|
||||
"""Find task by partial ID (supports short IDs)."""
|
||||
# First try exact match
|
||||
task = task_manager.find_task(partial_id)
|
||||
if task:
|
||||
return task
|
||||
|
||||
# Try partial match
|
||||
matching_tasks = [task for task in task_manager.tasks if task.id.startswith(partial_id)]
|
||||
|
||||
if len(matching_tasks) == 1:
|
||||
return matching_tasks[0]
|
||||
elif len(matching_tasks) > 1:
|
||||
print(f"Error: Multiple tasks match ID '{partial_id}':")
|
||||
for task in matching_tasks:
|
||||
print(f" {task.id[:8]} - {task.title}")
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
"""Main application entry point."""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.command:
|
||||
parser.print_help()
|
||||
return
|
||||
|
||||
task_manager = TaskManager()
|
||||
|
||||
try:
|
||||
if args.command == 'add':
|
||||
cmd_add(args, task_manager)
|
||||
elif args.command == 'list':
|
||||
cmd_list(args, task_manager)
|
||||
elif args.command == 'complete':
|
||||
# Handle partial ID matching
|
||||
task = find_task_by_partial_id(task_manager, args.id)
|
||||
if task:
|
||||
if task_manager.complete_task(task.id):
|
||||
print(f"Completed task: {task.title}")
|
||||
else:
|
||||
print(f"Error: Could not complete task '{task.id}'.")
|
||||
else:
|
||||
print(f"Error: Task with ID '{args.id}' not found.")
|
||||
elif args.command == 'delete':
|
||||
# Handle partial ID matching
|
||||
task = find_task_by_partial_id(task_manager, args.id)
|
||||
if task:
|
||||
task_title = task.title
|
||||
if task_manager.delete_task(task.id):
|
||||
print(f"Deleted task: {task_title}")
|
||||
else:
|
||||
print(f"Error: Could not delete task '{task.id}'.")
|
||||
else:
|
||||
print(f"Error: Task with ID '{args.id}' not found.")
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nOperation cancelled.")
|
||||
except Exception as e:
|
||||
print(f"An unexpected error occurred: {e}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,7 @@
|
||||
"""CSV Data Analysis Tool
|
||||
|
||||
A comprehensive tool for loading, analyzing, and reporting on CSV data.
|
||||
"""
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__author__ = "CSV Analysis Tool"
|
||||
@@ -0,0 +1,246 @@
|
||||
"""Statistical analysis, outlier detection, and correlation computation."""
|
||||
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
from scipy import stats
|
||||
from typing import Dict, List, Any, Tuple, Optional
|
||||
import warnings
|
||||
|
||||
warnings.filterwarnings('ignore', category=RuntimeWarning)
|
||||
|
||||
|
||||
class DataAnalyzer:
|
||||
"""Performs comprehensive statistical analysis on DataFrame."""
|
||||
|
||||
def __init__(self, df: pd.DataFrame):
|
||||
self.df = df
|
||||
self.numeric_columns = self._get_numeric_columns()
|
||||
self.categorical_columns = self._get_categorical_columns()
|
||||
self.datetime_columns = self._get_datetime_columns()
|
||||
|
||||
def _get_numeric_columns(self) -> List[str]:
|
||||
"""Get list of numeric columns."""
|
||||
return [col for col in self.df.columns if pd.api.types.is_numeric_dtype(self.df[col])]
|
||||
|
||||
def _get_categorical_columns(self) -> List[str]:
|
||||
"""Get list of categorical/string columns."""
|
||||
return [col for col in self.df.columns
|
||||
if pd.api.types.is_string_dtype(self.df[col]) or
|
||||
pd.api.types.is_categorical_dtype(self.df[col]) or
|
||||
self.df[col].dtype == 'object']
|
||||
|
||||
def _get_datetime_columns(self) -> List[str]:
|
||||
"""Get list of datetime columns."""
|
||||
return [col for col in self.df.columns if pd.api.types.is_datetime64_any_dtype(self.df[col])]
|
||||
|
||||
def compute_basic_info(self) -> Dict[str, Any]:
|
||||
"""Compute basic information about the dataset."""
|
||||
return {
|
||||
'total_rows': len(self.df),
|
||||
'total_columns': len(self.df.columns),
|
||||
'numeric_columns': len(self.numeric_columns),
|
||||
'categorical_columns': len(self.categorical_columns),
|
||||
'datetime_columns': len(self.datetime_columns),
|
||||
'memory_usage_mb': round(self.df.memory_usage(deep=True).sum() / (1024 * 1024), 2),
|
||||
'missing_values_total': self.df.isnull().sum().sum(),
|
||||
'duplicate_rows': self.df.duplicated().sum()
|
||||
}
|
||||
|
||||
def compute_summary_statistics(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Compute comprehensive summary statistics for all columns."""
|
||||
stats_dict = {}
|
||||
|
||||
# Numeric columns
|
||||
for col in self.numeric_columns:
|
||||
series = self.df[col].dropna()
|
||||
if len(series) == 0:
|
||||
stats_dict[col] = {'type': 'numeric', 'error': 'No non-null values'}
|
||||
continue
|
||||
|
||||
try:
|
||||
# Calculate mode safely
|
||||
mode_result = stats.mode(series, keepdims=False, nan_policy='omit')
|
||||
mode_value = mode_result.mode if hasattr(mode_result, 'mode') else mode_result[0]
|
||||
|
||||
# Handle case where mode might be an array
|
||||
if isinstance(mode_value, np.ndarray):
|
||||
mode_value = mode_value[0] if len(mode_value) > 0 else np.nan
|
||||
|
||||
quartiles = series.quantile([0.25, 0.5, 0.75])
|
||||
|
||||
stats_dict[col] = {
|
||||
'type': 'numeric',
|
||||
'count': len(series),
|
||||
'missing': self.df[col].isnull().sum(),
|
||||
'mean': round(series.mean(), 4),
|
||||
'median': round(series.median(), 4),
|
||||
'mode': round(float(mode_value), 4) if not pd.isna(mode_value) else None,
|
||||
'std_dev': round(series.std(), 4),
|
||||
'variance': round(series.var(), 4),
|
||||
'min': round(series.min(), 4),
|
||||
'max': round(series.max(), 4),
|
||||
'q1': round(quartiles[0.25], 4),
|
||||
'q3': round(quartiles[0.75], 4),
|
||||
'iqr': round(quartiles[0.75] - quartiles[0.25], 4),
|
||||
'skewness': round(series.skew(), 4),
|
||||
'kurtosis': round(series.kurtosis(), 4)
|
||||
}
|
||||
except Exception as e:
|
||||
stats_dict[col] = {'type': 'numeric', 'error': str(e)}
|
||||
|
||||
# Categorical columns
|
||||
for col in self.categorical_columns:
|
||||
series = self.df[col].dropna()
|
||||
if len(series) == 0:
|
||||
stats_dict[col] = {'type': 'categorical', 'error': 'No non-null values'}
|
||||
continue
|
||||
|
||||
try:
|
||||
value_counts = series.value_counts()
|
||||
stats_dict[col] = {
|
||||
'type': 'categorical',
|
||||
'count': len(series),
|
||||
'missing': self.df[col].isnull().sum(),
|
||||
'unique_values': series.nunique(),
|
||||
'most_frequent': value_counts.index[0] if len(value_counts) > 0 else None,
|
||||
'most_frequent_count': value_counts.iloc[0] if len(value_counts) > 0 else 0,
|
||||
'least_frequent': value_counts.index[-1] if len(value_counts) > 0 else None,
|
||||
'least_frequent_count': value_counts.iloc[-1] if len(value_counts) > 0 else 0,
|
||||
'top_5_values': dict(value_counts.head().items())
|
||||
}
|
||||
except Exception as e:
|
||||
stats_dict[col] = {'type': 'categorical', 'error': str(e)}
|
||||
|
||||
# Datetime columns
|
||||
for col in self.datetime_columns:
|
||||
series = self.df[col].dropna()
|
||||
if len(series) == 0:
|
||||
stats_dict[col] = {'type': 'datetime', 'error': 'No non-null values'}
|
||||
continue
|
||||
|
||||
try:
|
||||
stats_dict[col] = {
|
||||
'type': 'datetime',
|
||||
'count': len(series),
|
||||
'missing': self.df[col].isnull().sum(),
|
||||
'min_date': series.min().strftime('%Y-%m-%d %H:%M:%S'),
|
||||
'max_date': series.max().strftime('%Y-%m-%d %H:%M:%S'),
|
||||
'date_range_days': (series.max() - series.min()).days,
|
||||
'unique_dates': series.nunique()
|
||||
}
|
||||
except Exception as e:
|
||||
stats_dict[col] = {'type': 'datetime', 'error': str(e)}
|
||||
|
||||
return stats_dict
|
||||
|
||||
def detect_outliers_iqr(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Detect outliers using Interquartile Range (IQR) method."""
|
||||
outliers_dict = {}
|
||||
|
||||
for col in self.numeric_columns:
|
||||
series = self.df[col].dropna()
|
||||
if len(series) < 4: # Need at least 4 values for quartiles
|
||||
outliers_dict[col] = {'error': 'Insufficient data for outlier detection'}
|
||||
continue
|
||||
|
||||
try:
|
||||
q1 = series.quantile(0.25)
|
||||
q3 = series.quantile(0.75)
|
||||
iqr = q3 - q1
|
||||
|
||||
lower_bound = q1 - 1.5 * iqr
|
||||
upper_bound = q3 + 1.5 * iqr
|
||||
|
||||
outliers = series[(series < lower_bound) | (series > upper_bound)]
|
||||
outlier_indices = self.df[col][(self.df[col] < lower_bound) | (self.df[col] > upper_bound)].index.tolist()
|
||||
|
||||
outliers_dict[col] = {
|
||||
'lower_bound': round(lower_bound, 4),
|
||||
'upper_bound': round(upper_bound, 4),
|
||||
'outlier_count': len(outliers),
|
||||
'outlier_percentage': round((len(outliers) / len(series)) * 100, 2),
|
||||
'outlier_values': [round(x, 4) for x in sorted(outliers.values)],
|
||||
'outlier_indices': outlier_indices[:20] # Limit to first 20 indices
|
||||
}
|
||||
except Exception as e:
|
||||
outliers_dict[col] = {'error': str(e)}
|
||||
|
||||
return outliers_dict
|
||||
|
||||
def compute_correlation_matrix(self) -> Optional[Dict[str, Any]]:
|
||||
"""Compute correlation matrix for numeric columns."""
|
||||
if len(self.numeric_columns) < 2:
|
||||
return {'error': 'Need at least 2 numeric columns for correlation analysis'}
|
||||
|
||||
try:
|
||||
numeric_df = self.df[self.numeric_columns].select_dtypes(include=[np.number])
|
||||
|
||||
# Remove columns with no variance
|
||||
numeric_df = numeric_df.loc[:, numeric_df.var() != 0]
|
||||
|
||||
if len(numeric_df.columns) < 2:
|
||||
return {'error': 'Need at least 2 numeric columns with variance for correlation'}
|
||||
|
||||
# Compute correlations
|
||||
pearson_corr = numeric_df.corr(method='pearson')
|
||||
spearman_corr = numeric_df.corr(method='spearman')
|
||||
|
||||
# Find strong correlations (> 0.7 or < -0.7)
|
||||
strong_correlations = []
|
||||
for i in range(len(pearson_corr.columns)):
|
||||
for j in range(i + 1, len(pearson_corr.columns)):
|
||||
col1, col2 = pearson_corr.columns[i], pearson_corr.columns[j]
|
||||
pearson_val = pearson_corr.iloc[i, j]
|
||||
spearman_val = spearman_corr.iloc[i, j]
|
||||
|
||||
if abs(pearson_val) > 0.7 and not pd.isna(pearson_val):
|
||||
strong_correlations.append({
|
||||
'column1': col1,
|
||||
'column2': col2,
|
||||
'pearson': round(pearson_val, 4),
|
||||
'spearman': round(spearman_val, 4)
|
||||
})
|
||||
|
||||
return {
|
||||
'pearson_correlation': pearson_corr.round(4).to_dict(),
|
||||
'spearman_correlation': spearman_corr.round(4).to_dict(),
|
||||
'strong_correlations': strong_correlations,
|
||||
'columns_analyzed': list(numeric_df.columns)
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {'error': f'Error computing correlations: {str(e)}'}
|
||||
|
||||
def analyze_missing_patterns(self) -> Dict[str, Any]:
|
||||
"""Analyze missing data patterns."""
|
||||
missing_data = {}
|
||||
|
||||
# Missing values per column
|
||||
missing_per_column = self.df.isnull().sum()
|
||||
missing_percentages = (missing_per_column / len(self.df) * 100).round(2)
|
||||
|
||||
# Columns with missing data
|
||||
columns_with_missing = missing_per_column[missing_per_column > 0].to_dict()
|
||||
|
||||
# Missing data patterns
|
||||
missing_combinations = self.df.isnull().value_counts().head(10)
|
||||
|
||||
return {
|
||||
'total_missing_values': missing_per_column.sum(),
|
||||
'missing_percentage_overall': round((missing_per_column.sum() / (len(self.df) * len(self.df.columns))) * 100, 2),
|
||||
'columns_with_missing': dict(zip(columns_with_missing.keys(),
|
||||
[{'count': int(count), 'percentage': float(missing_percentages[col])}
|
||||
for col, count in columns_with_missing.items()])),
|
||||
'complete_rows': len(self.df.dropna()),
|
||||
'complete_rows_percentage': round((len(self.df.dropna()) / len(self.df)) * 100, 2)
|
||||
}
|
||||
|
||||
def run_full_analysis(self) -> Dict[str, Any]:
|
||||
"""Run complete analysis and return all results."""
|
||||
return {
|
||||
'basic_info': self.compute_basic_info(),
|
||||
'summary_statistics': self.compute_summary_statistics(),
|
||||
'outliers': self.detect_outliers_iqr(),
|
||||
'correlations': self.compute_correlation_matrix(),
|
||||
'missing_data': self.analyze_missing_patterns()
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Command-line interface for CSV data analysis tool."""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
from loader import CSVLoader
|
||||
from analyzer import DataAnalyzer
|
||||
from reporter import ReportFormatter
|
||||
|
||||
|
||||
def parse_arguments():
|
||||
"""Parse command-line arguments."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="CSV Data Analysis Tool - Comprehensive statistical analysis of CSV files",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python cli.py --file data.csv
|
||||
python cli.py --file data.csv --output report.txt
|
||||
python cli.py --file data.csv --columns age,salary,score
|
||||
python cli.py --file data.csv --filter "age > 25" "salary >= 50000"
|
||||
python cli.py --file data.csv --columns age,salary --filter "age > 18" --output analysis.txt
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--file', '-f',
|
||||
type=str,
|
||||
required=True,
|
||||
help='Path to the CSV file to analyze'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--columns', '-c',
|
||||
type=str,
|
||||
help='Comma-separated list of columns to analyze (default: all columns)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--filter',
|
||||
type=str,
|
||||
nargs='+',
|
||||
help='Filter conditions (e.g., "age > 25" "name contains John"). Supported operators: ==, !=, >, >=, <, <=, contains'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--output', '-o',
|
||||
type=str,
|
||||
help='Output file path for the report (default: print to console)'
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def validate_arguments(args):
|
||||
"""Validate command-line arguments."""
|
||||
# Check if file exists
|
||||
if not os.path.exists(args.file):
|
||||
print(f"Error: File '{args.file}' not found.")
|
||||
sys.exit(1)
|
||||
|
||||
# Validate file extension
|
||||
if not args.file.lower().endswith(('.csv', '.txt')):
|
||||
print("Warning: File does not have a .csv extension. Proceeding anyway...")
|
||||
|
||||
# Parse columns if provided
|
||||
columns = None
|
||||
if args.columns:
|
||||
columns = [col.strip() for col in args.columns.split(',') if col.strip()]
|
||||
if not columns:
|
||||
print("Error: Invalid columns specification.")
|
||||
sys.exit(1)
|
||||
|
||||
# Validate output path
|
||||
if args.output:
|
||||
output_dir = os.path.dirname(os.path.abspath(args.output))
|
||||
if not os.path.exists(output_dir):
|
||||
print(f"Error: Output directory '{output_dir}' does not exist.")
|
||||
sys.exit(1)
|
||||
|
||||
return columns
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to orchestrate the analysis pipeline."""
|
||||
try:
|
||||
# Parse and validate arguments
|
||||
args = parse_arguments()
|
||||
columns = validate_arguments(args)
|
||||
|
||||
print("CSV Data Analysis Tool")
|
||||
print("=" * 50)
|
||||
|
||||
# Initialize loader
|
||||
loader = CSVLoader()
|
||||
|
||||
# Load CSV file
|
||||
print(f"Loading file: {args.file}")
|
||||
try:
|
||||
df = loader.load_csv(args.file, columns=columns)
|
||||
except Exception as e:
|
||||
print(f"Error loading CSV file: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Apply filters if specified
|
||||
if args.filter:
|
||||
print(f"Applying filters: {args.filter}")
|
||||
try:
|
||||
df = loader.apply_filters(df, args.filter)
|
||||
if len(df) == 0:
|
||||
print("Warning: All rows were filtered out. No data to analyze.")
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"Error applying filters: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Initialize analyzer
|
||||
print("Performing statistical analysis...")
|
||||
analyzer = DataAnalyzer(df)
|
||||
|
||||
# Run analysis
|
||||
try:
|
||||
analysis_results = analyzer.run_full_analysis()
|
||||
except Exception as e:
|
||||
print(f"Error during analysis: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Generate report
|
||||
print("Generating report...")
|
||||
try:
|
||||
formatter = ReportFormatter()
|
||||
report_text = formatter.generate_report(
|
||||
analysis_results=analysis_results,
|
||||
file_path=args.file,
|
||||
filters_applied=args.filter
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error generating report: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Output report
|
||||
if args.output:
|
||||
try:
|
||||
formatter.save_report(report_text, args.output)
|
||||
print(f"\nAnalysis complete! Report saved to: {args.output}")
|
||||
except Exception as e:
|
||||
print(f"Error saving report: {e}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("\n" + "=" * 50)
|
||||
print(report_text)
|
||||
|
||||
# Print summary
|
||||
basic_info = analysis_results.get('basic_info', {})
|
||||
print(f"\nSummary:")
|
||||
print(f" Rows analyzed: {basic_info.get('total_rows', 'N/A'):,}")
|
||||
print(f" Columns analyzed: {basic_info.get('total_columns', 'N/A')}")
|
||||
|
||||
if args.filter:
|
||||
print(f" Filters applied: {len(args.filter)}")
|
||||
|
||||
outliers = analysis_results.get('outliers', {})
|
||||
total_outliers = sum(info.get('outlier_count', 0) for info in outliers.values()
|
||||
if isinstance(info, dict) and 'outlier_count' in info)
|
||||
if total_outliers > 0:
|
||||
print(f" Total outliers detected: {total_outliers}")
|
||||
|
||||
correlations = analysis_results.get('correlations', {})
|
||||
if isinstance(correlations, dict) and 'strong_correlations' in correlations:
|
||||
strong_corr_count = len(correlations['strong_correlations'])
|
||||
if strong_corr_count > 0:
|
||||
print(f" Strong correlations found: {strong_corr_count}")
|
||||
|
||||
print("\nAnalysis completed successfully!")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nOperation cancelled by user.")
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"Unexpected error: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,192 @@
|
||||
"""CSV file loading with automatic type detection and encoding handling."""
|
||||
|
||||
import pandas as pd
|
||||
import chardet
|
||||
import numpy as np
|
||||
from typing import Optional, Dict, Any
|
||||
import os
|
||||
|
||||
|
||||
class CSVLoader:
|
||||
"""Handles CSV file loading with automatic encoding detection and type inference."""
|
||||
|
||||
def __init__(self):
|
||||
self.encoding = None
|
||||
self.delimiter = None
|
||||
|
||||
def detect_encoding(self, file_path: str) -> str:
|
||||
"""Detect the encoding of a CSV file."""
|
||||
try:
|
||||
with open(file_path, 'rb') as file:
|
||||
raw_data = file.read(10000) # Read first 10KB
|
||||
result = chardet.detect(raw_data)
|
||||
encoding = result.get('encoding', 'utf-8')
|
||||
|
||||
if encoding is None or result.get('confidence', 0) < 0.7:
|
||||
encoding = 'utf-8'
|
||||
|
||||
self.encoding = encoding
|
||||
return encoding
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not detect encoding, using utf-8. Error: {e}")
|
||||
self.encoding = 'utf-8'
|
||||
return 'utf-8'
|
||||
|
||||
def detect_delimiter(self, file_path: str, encoding: str) -> str:
|
||||
"""Detect the delimiter used in the CSV file."""
|
||||
delimiters = [',', ';', '\t', '|']
|
||||
delimiter_counts = {}
|
||||
|
||||
try:
|
||||
with open(file_path, 'r', encoding=encoding) as file:
|
||||
first_line = file.readline()
|
||||
for delimiter in delimiters:
|
||||
delimiter_counts[delimiter] = first_line.count(delimiter)
|
||||
|
||||
# Choose delimiter with highest count
|
||||
best_delimiter = max(delimiter_counts, key=delimiter_counts.get)
|
||||
self.delimiter = best_delimiter if delimiter_counts[best_delimiter] > 0 else ','
|
||||
return self.delimiter
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not detect delimiter, using comma. Error: {e}")
|
||||
self.delimiter = ','
|
||||
return ','
|
||||
|
||||
def infer_column_types(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
"""Automatically infer and convert column types."""
|
||||
df_copy = df.copy()
|
||||
|
||||
for column in df_copy.columns:
|
||||
# Skip if already numeric
|
||||
if pd.api.types.is_numeric_dtype(df_copy[column]):
|
||||
continue
|
||||
|
||||
# Try to convert to numeric first
|
||||
numeric_series = pd.to_numeric(df_copy[column], errors='coerce')
|
||||
if not numeric_series.isna().all():
|
||||
# If more than 80% of values can be converted to numeric, treat as numeric
|
||||
non_null_count = df_copy[column].count()
|
||||
numeric_count = numeric_series.count()
|
||||
if non_null_count > 0 and (numeric_count / non_null_count) >= 0.8:
|
||||
df_copy[column] = numeric_series
|
||||
continue
|
||||
|
||||
# Try to convert to datetime
|
||||
try:
|
||||
datetime_series = pd.to_datetime(df_copy[column], errors='coerce', infer_datetime_format=True)
|
||||
if not datetime_series.isna().all():
|
||||
non_null_count = df_copy[column].count()
|
||||
datetime_count = datetime_series.count()
|
||||
if non_null_count > 0 and (datetime_count / non_null_count) >= 0.8:
|
||||
df_copy[column] = datetime_series
|
||||
continue
|
||||
except:
|
||||
pass
|
||||
|
||||
# Keep as categorical/string
|
||||
df_copy[column] = df_copy[column].astype('string')
|
||||
|
||||
return df_copy
|
||||
|
||||
def load_csv(self, file_path: str, columns: Optional[list] = None) -> pd.DataFrame:
|
||||
"""Load CSV file with automatic encoding detection and type inference."""
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
# Detect encoding and delimiter
|
||||
encoding = self.detect_encoding(file_path)
|
||||
delimiter = self.detect_delimiter(file_path, encoding)
|
||||
|
||||
try:
|
||||
# Load the CSV
|
||||
df = pd.read_csv(
|
||||
file_path,
|
||||
encoding=encoding,
|
||||
delimiter=delimiter,
|
||||
low_memory=False
|
||||
)
|
||||
|
||||
# Filter columns if specified
|
||||
if columns:
|
||||
available_columns = [col for col in columns if col in df.columns]
|
||||
missing_columns = [col for col in columns if col not in df.columns]
|
||||
|
||||
if missing_columns:
|
||||
print(f"Warning: Columns not found in file: {missing_columns}")
|
||||
|
||||
if not available_columns:
|
||||
raise ValueError("None of the specified columns were found in the file")
|
||||
|
||||
df = df[available_columns]
|
||||
|
||||
# Infer and convert column types
|
||||
df = self.infer_column_types(df)
|
||||
|
||||
print(f"Successfully loaded CSV with {len(df)} rows and {len(df.columns)} columns")
|
||||
print(f"Encoding: {encoding}, Delimiter: '{delimiter}'")
|
||||
|
||||
return df
|
||||
|
||||
except pd.errors.EmptyDataError:
|
||||
raise ValueError("The CSV file is empty")
|
||||
except pd.errors.ParserError as e:
|
||||
raise ValueError(f"Error parsing CSV file: {e}")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error loading CSV file: {e}")
|
||||
|
||||
def apply_filters(self, df: pd.DataFrame, filter_conditions: list) -> pd.DataFrame:
|
||||
"""Apply filtering conditions to the DataFrame."""
|
||||
if not filter_conditions:
|
||||
return df
|
||||
|
||||
filtered_df = df.copy()
|
||||
|
||||
for condition in filter_conditions:
|
||||
try:
|
||||
# Parse condition (format: column_name operator value)
|
||||
parts = condition.split(' ', 2)
|
||||
if len(parts) != 3:
|
||||
print(f"Warning: Invalid filter format '{condition}'. Expected: 'column operator value'")
|
||||
continue
|
||||
|
||||
column, operator, value = parts
|
||||
|
||||
if column not in filtered_df.columns:
|
||||
print(f"Warning: Column '{column}' not found. Skipping filter.")
|
||||
continue
|
||||
|
||||
# Convert value to appropriate type
|
||||
if pd.api.types.is_numeric_dtype(filtered_df[column]):
|
||||
try:
|
||||
value = float(value)
|
||||
except ValueError:
|
||||
print(f"Warning: Cannot convert '{value}' to numeric for column '{column}'")
|
||||
continue
|
||||
|
||||
# Apply filter based on operator
|
||||
if operator == '==':
|
||||
mask = filtered_df[column] == value
|
||||
elif operator == '!=':
|
||||
mask = filtered_df[column] != value
|
||||
elif operator == '>':
|
||||
mask = filtered_df[column] > value
|
||||
elif operator == '>=':
|
||||
mask = filtered_df[column] >= value
|
||||
elif operator == '<':
|
||||
mask = filtered_df[column] < value
|
||||
elif operator == '<=':
|
||||
mask = filtered_df[column] <= value
|
||||
elif operator == 'contains':
|
||||
mask = filtered_df[column].astype(str).str.contains(str(value), na=False)
|
||||
else:
|
||||
print(f"Warning: Unsupported operator '{operator}'. Supported: ==, !=, >, >=, <, <=, contains")
|
||||
continue
|
||||
|
||||
filtered_df = filtered_df[mask]
|
||||
|
||||
except Exception as e:
|
||||
print(f"Warning: Error applying filter '{condition}': {e}")
|
||||
continue
|
||||
|
||||
print(f"Applied {len(filter_conditions)} filters. Rows remaining: {len(filtered_df)}")
|
||||
return filtered_df
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Format analysis results into readable text reports."""
|
||||
|
||||
from typing import Dict, Any, List
|
||||
import pandas as pd
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class ReportFormatter:
|
||||
"""Formats analysis results into comprehensive text reports."""
|
||||
|
||||
def __init__(self):
|
||||
self.report_lines = []
|
||||
|
||||
def _add_line(self, text: str = "", level: int = 0):
|
||||
"""Add a line to the report with proper indentation."""
|
||||
indent = " " * level
|
||||
self.report_lines.append(f"{indent}{text}")
|
||||
|
||||
def _add_separator(self, char: str = "=", length: int = 80):
|
||||
"""Add a separator line."""
|
||||
self.report_lines.append(char * length)
|
||||
|
||||
def _add_header(self, title: str, level: int = 1):
|
||||
"""Add a formatted header."""
|
||||
if level == 1:
|
||||
self._add_separator("=")
|
||||
self._add_line(f" {title.upper()} ")
|
||||
self._add_separator("=")
|
||||
elif level == 2:
|
||||
self._add_separator("-")
|
||||
self._add_line(f"{title}")
|
||||
self._add_separator("-")
|
||||
else:
|
||||
self._add_line(f"\n{title}:")
|
||||
|
||||
def format_basic_info(self, basic_info: Dict[str, Any]):
|
||||
"""Format basic dataset information."""
|
||||
self._add_header("Dataset Overview", level=1)
|
||||
self._add_line(f"Total Rows: {basic_info.get('total_rows', 'N/A'):,}")
|
||||
self._add_line(f"Total Columns: {basic_info.get('total_columns', 'N/A')}")
|
||||
self._add_line(f" - Numeric Columns: {basic_info.get('numeric_columns', 'N/A')}")
|
||||
self._add_line(f" - Categorical Columns: {basic_info.get('categorical_columns', 'N/A')}")
|
||||
self._add_line(f" - Datetime Columns: {basic_info.get('datetime_columns', 'N/A')}")
|
||||
self._add_line(f"Memory Usage: {basic_info.get('memory_usage_mb', 'N/A')} MB")
|
||||
self._add_line(f"Missing Values: {basic_info.get('missing_values_total', 'N/A'):,}")
|
||||
self._add_line(f"Duplicate Rows: {basic_info.get('duplicate_rows', 'N/A'):,}")
|
||||
self._add_line()
|
||||
|
||||
def format_summary_statistics(self, summary_stats: Dict[str, Dict[str, Any]]):
|
||||
"""Format summary statistics for all columns."""
|
||||
self._add_header("Summary Statistics", level=1)
|
||||
|
||||
# Group by column type
|
||||
numeric_cols = {k: v for k, v in summary_stats.items() if v.get('type') == 'numeric'}
|
||||
categorical_cols = {k: v for k, v in summary_stats.items() if v.get('type') == 'categorical'}
|
||||
datetime_cols = {k: v for k, v in summary_stats.items() if v.get('type') == 'datetime'}
|
||||
|
||||
# Numeric columns
|
||||
if numeric_cols:
|
||||
self._add_header("Numeric Columns", level=2)
|
||||
for col_name, stats in numeric_cols.items():
|
||||
if 'error' in stats:
|
||||
self._add_line(f"{col_name}: Error - {stats['error']}")
|
||||
continue
|
||||
|
||||
self._add_line(f"\n{col_name}:")
|
||||
self._add_line(f"Count: {stats.get('count', 'N/A'):,}", 1)
|
||||
self._add_line(f"Missing: {stats.get('missing', 'N/A'):,}", 1)
|
||||
self._add_line(f"Mean: {stats.get('mean', 'N/A')}", 1)
|
||||
self._add_line(f"Median: {stats.get('median', 'N/A')}", 1)
|
||||
if stats.get('mode') is not None:
|
||||
self._add_line(f"Mode: {stats.get('mode', 'N/A')}", 1)
|
||||
self._add_line(f"Std Dev: {stats.get('std_dev', 'N/A')}", 1)
|
||||
self._add_line(f"Min: {stats.get('min', 'N/A')}", 1)
|
||||
self._add_line(f"Q1: {stats.get('q1', 'N/A')}", 1)
|
||||
self._add_line(f"Q3: {stats.get('q3', 'N/A')}", 1)
|
||||
self._add_line(f"Max: {stats.get('max', 'N/A')}", 1)
|
||||
self._add_line(f"IQR: {stats.get('iqr', 'N/A')}", 1)
|
||||
self._add_line(f"Skewness: {stats.get('skewness', 'N/A')}", 1)
|
||||
self._add_line(f"Kurtosis: {stats.get('kurtosis', 'N/A')}", 1)
|
||||
|
||||
# Categorical columns
|
||||
if categorical_cols:
|
||||
self._add_header("Categorical Columns", level=2)
|
||||
for col_name, stats in categorical_cols.items():
|
||||
if 'error' in stats:
|
||||
self._add_line(f"{col_name}: Error - {stats['error']}")
|
||||
continue
|
||||
|
||||
self._add_line(f"\n{col_name}:")
|
||||
self._add_line(f"Count: {stats.get('count', 'N/A'):,}", 1)
|
||||
self._add_line(f"Missing: {stats.get('missing', 'N/A'):,}", 1)
|
||||
self._add_line(f"Unique Values: {stats.get('unique_values', 'N/A'):,}", 1)
|
||||
self._add_line(f"Most Frequent: '{stats.get('most_frequent', 'N/A')}' ({stats.get('most_frequent_count', 'N/A')} times)", 1)
|
||||
|
||||
if stats.get('top_5_values'):
|
||||
self._add_line("Top 5 Values:", 1)
|
||||
for value, count in list(stats['top_5_values'].items())[:5]:
|
||||
self._add_line(f" '{value}': {count}", 2)
|
||||
|
||||
# Datetime columns
|
||||
if datetime_cols:
|
||||
self._add_header("Datetime Columns", level=2)
|
||||
for col_name, stats in datetime_cols.items():
|
||||
if 'error' in stats:
|
||||
self._add_line(f"{col_name}: Error - {stats['error']}")
|
||||
continue
|
||||
|
||||
self._add_line(f"\n{col_name}:")
|
||||
self._add_line(f"Count: {stats.get('count', 'N/A'):,}", 1)
|
||||
self._add_line(f"Missing: {stats.get('missing', 'N/A'):,}", 1)
|
||||
self._add_line(f"Date Range: {stats.get('min_date', 'N/A')} to {stats.get('max_date', 'N/A')}", 1)
|
||||
self._add_line(f"Range (days): {stats.get('date_range_days', 'N/A'):,}", 1)
|
||||
self._add_line(f"Unique Dates: {stats.get('unique_dates', 'N/A'):,}", 1)
|
||||
|
||||
self._add_line()
|
||||
|
||||
def format_outliers(self, outliers: Dict[str, Dict[str, Any]]):
|
||||
"""Format outlier detection results."""
|
||||
self._add_header("Outlier Detection (IQR Method)", level=1)
|
||||
|
||||
if not outliers:
|
||||
self._add_line("No numeric columns available for outlier detection.")
|
||||
return
|
||||
|
||||
has_outliers = False
|
||||
for col_name, outlier_info in outliers.items():
|
||||
if 'error' in outlier_info:
|
||||
self._add_line(f"{col_name}: {outlier_info['error']}")
|
||||
continue
|
||||
|
||||
outlier_count = outlier_info.get('outlier_count', 0)
|
||||
if outlier_count > 0:
|
||||
has_outliers = True
|
||||
self._add_line(f"\n{col_name}:")
|
||||
self._add_line(f"Outliers Found: {outlier_count} ({outlier_info.get('outlier_percentage', 0)}% of data)", 1)
|
||||
self._add_line(f"Valid Range: [{outlier_info.get('lower_bound', 'N/A')}, {outlier_info.get('upper_bound', 'N/A')}]", 1)
|
||||
|
||||
outlier_values = outlier_info.get('outlier_values', [])
|
||||
if outlier_values:
|
||||
display_values = outlier_values[:10] # Show first 10 outliers
|
||||
self._add_line(f"Sample Outliers: {display_values}", 1)
|
||||
if len(outlier_values) > 10:
|
||||
self._add_line(f"... and {len(outlier_values) - 10} more", 1)
|
||||
|
||||
if not has_outliers:
|
||||
self._add_line("No outliers detected in any numeric columns.")
|
||||
|
||||
self._add_line()
|
||||
|
||||
def format_correlations(self, correlations: Dict[str, Any]):
|
||||
"""Format correlation analysis results."""
|
||||
self._add_header("Correlation Analysis", level=1)
|
||||
|
||||
if 'error' in correlations:
|
||||
self._add_line(f"Error: {correlations['error']}")
|
||||
self._add_line()
|
||||
return
|
||||
|
||||
columns_analyzed = correlations.get('columns_analyzed', [])
|
||||
self._add_line(f"Columns analyzed: {', '.join(columns_analyzed)}")
|
||||
|
||||
# Strong correlations
|
||||
strong_corr = correlations.get('strong_correlations', [])
|
||||
if strong_corr:
|
||||
self._add_header("Strong Correlations (|r| > 0.7)", level=3)
|
||||
for corr in strong_corr:
|
||||
self._add_line(f"{corr['column1']} <-> {corr['column2']}:")
|
||||
self._add_line(f" Pearson: {corr['pearson']}", 1)
|
||||
self._add_line(f" Spearman: {corr['spearman']}", 1)
|
||||
else:
|
||||
self._add_line("\nNo strong correlations (|r| > 0.7) found.")
|
||||
|
||||
# Correlation matrix (Pearson)
|
||||
pearson_corr = correlations.get('pearson_correlation', {})
|
||||
if pearson_corr and len(columns_analyzed) <= 10: # Only show matrix for smaller datasets
|
||||
self._add_header("Pearson Correlation Matrix", level=3)
|
||||
|
||||
# Create a formatted correlation matrix
|
||||
matrix_lines = []
|
||||
col_names = list(pearson_corr.keys())
|
||||
|
||||
# Header line
|
||||
header = "Variable".ljust(15) + "".join([col[:8].rjust(10) for col in col_names])
|
||||
matrix_lines.append(header)
|
||||
matrix_lines.append("-" * len(header))
|
||||
|
||||
# Matrix rows
|
||||
for row_name in col_names:
|
||||
row_data = pearson_corr[row_name]
|
||||
row_line = row_name[:14].ljust(15)
|
||||
for col_name in col_names:
|
||||
corr_val = row_data.get(col_name, 0)
|
||||
row_line += f"{corr_val:8.3f} "
|
||||
matrix_lines.append(row_line)
|
||||
|
||||
for line in matrix_lines:
|
||||
self._add_line(line)
|
||||
|
||||
self._add_line()
|
||||
|
||||
def format_missing_data(self, missing_data: Dict[str, Any]):
|
||||
"""Format missing data analysis results."""
|
||||
self._add_header("Missing Data Analysis", level=1)
|
||||
|
||||
self._add_line(f"Total Missing Values: {missing_data.get('total_missing_values', 'N/A'):,}")
|
||||
self._add_line(f"Overall Missing Percentage: {missing_data.get('missing_percentage_overall', 'N/A')}%")
|
||||
self._add_line(f"Complete Rows: {missing_data.get('complete_rows', 'N/A'):,} ({missing_data.get('complete_rows_percentage', 'N/A')}%)")
|
||||
|
||||
columns_with_missing = missing_data.get('columns_with_missing', {})
|
||||
if columns_with_missing:
|
||||
self._add_header("Columns with Missing Data", level=3)
|
||||
for col_name, info in columns_with_missing.items():
|
||||
self._add_line(f"{col_name}: {info['count']:,} missing ({info['percentage']}%)")
|
||||
else:
|
||||
self._add_line("\nNo missing data found in any columns.")
|
||||
|
||||
self._add_line()
|
||||
|
||||
def generate_report(self, analysis_results: Dict[str, Any],
|
||||
file_path: str = None, filters_applied: List[str] = None) -> str:
|
||||
"""Generate complete formatted report."""
|
||||
self.report_lines = []
|
||||
|
||||
# Report header
|
||||
self._add_header("CSV DATA ANALYSIS REPORT", level=1)
|
||||
self._add_line(f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
if file_path:
|
||||
self._add_line(f"Source file: {file_path}")
|
||||
if filters_applied:
|
||||
self._add_line(f"Filters applied: {', '.join(filters_applied)}")
|
||||
self._add_line()
|
||||
|
||||
# Add each section
|
||||
if 'basic_info' in analysis_results:
|
||||
self.format_basic_info(analysis_results['basic_info'])
|
||||
|
||||
if 'summary_statistics' in analysis_results:
|
||||
self.format_summary_statistics(analysis_results['summary_statistics'])
|
||||
|
||||
if 'outliers' in analysis_results:
|
||||
self.format_outliers(analysis_results['outliers'])
|
||||
|
||||
if 'correlations' in analysis_results:
|
||||
self.format_correlations(analysis_results['correlations'])
|
||||
|
||||
if 'missing_data' in analysis_results:
|
||||
self.format_missing_data(analysis_results['missing_data'])
|
||||
|
||||
# Footer
|
||||
self._add_separator("=")
|
||||
self._add_line("End of Report")
|
||||
self._add_separator("=")
|
||||
|
||||
return "\n".join(self.report_lines)
|
||||
|
||||
def save_report(self, report_text: str, output_path: str):
|
||||
"""Save the report to a file."""
|
||||
try:
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
f.write(report_text)
|
||||
print(f"Report saved to: {output_path}")
|
||||
except Exception as e:
|
||||
raise IOError(f"Error saving report to {output_path}: {e}")
|
||||
@@ -0,0 +1,4 @@
|
||||
pandas>=1.3.0
|
||||
numpy>=1.20.0
|
||||
chardet>=4.0.0
|
||||
scipy>=1.7.0
|
||||
@@ -0,0 +1,19 @@
|
||||
"""
|
||||
Rate limiting middleware library for Flask APIs.
|
||||
"""
|
||||
|
||||
from .algorithms import TokenBucket, SlidingWindow, FixedWindow
|
||||
from .config import RateLimitConfig, LimitType
|
||||
from .middleware import rate_limit
|
||||
from .store import InMemoryStore
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__all__ = [
|
||||
"TokenBucket",
|
||||
"SlidingWindow",
|
||||
"FixedWindow",
|
||||
"RateLimitConfig",
|
||||
"LimitType",
|
||||
"rate_limit",
|
||||
"InMemoryStore"
|
||||
]
|
||||
@@ -0,0 +1,151 @@
|
||||
import time
|
||||
import math
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Tuple, Dict, Any, List
|
||||
from .store import store
|
||||
|
||||
|
||||
class RateLimitAlgorithm(ABC):
|
||||
"""Abstract base class for rate limiting algorithms."""
|
||||
|
||||
def __init__(self, limit: int, window: int):
|
||||
self.limit = limit
|
||||
self.window = window
|
||||
|
||||
@abstractmethod
|
||||
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
||||
"""
|
||||
Check if request is allowed.
|
||||
Returns (allowed, headers_dict).
|
||||
"""
|
||||
pass
|
||||
|
||||
def _get_key(self, identifier: str) -> str:
|
||||
"""Generate storage key for identifier."""
|
||||
return f"{self.__class__.__name__}:{identifier}"
|
||||
|
||||
|
||||
class TokenBucket(RateLimitAlgorithm):
|
||||
"""Token bucket rate limiting algorithm."""
|
||||
|
||||
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
||||
key = self._get_key(identifier)
|
||||
now = time.time()
|
||||
|
||||
# Get current bucket state
|
||||
bucket_data = store.get(key)
|
||||
|
||||
if bucket_data is None:
|
||||
# Initialize bucket
|
||||
bucket_data = {
|
||||
'tokens': self.limit,
|
||||
'last_refill': now
|
||||
}
|
||||
else:
|
||||
# Refill tokens based on elapsed time
|
||||
elapsed = now - bucket_data['last_refill']
|
||||
tokens_to_add = elapsed * (self.limit / self.window)
|
||||
bucket_data['tokens'] = min(self.limit, bucket_data['tokens'] + tokens_to_add)
|
||||
bucket_data['last_refill'] = now
|
||||
|
||||
# Check if request can be served
|
||||
if bucket_data['tokens'] >= 1:
|
||||
bucket_data['tokens'] -= 1
|
||||
allowed = True
|
||||
remaining = int(bucket_data['tokens'])
|
||||
else:
|
||||
allowed = False
|
||||
remaining = 0
|
||||
|
||||
# Store updated bucket state (TTL = 2 * window for safety)
|
||||
store.set(key, bucket_data, ttl=self.window * 2)
|
||||
|
||||
# Calculate reset time
|
||||
if bucket_data['tokens'] < self.limit:
|
||||
time_to_full = (self.limit - bucket_data['tokens']) / (self.limit / self.window)
|
||||
reset_time = int(now + time_to_full)
|
||||
else:
|
||||
reset_time = int(now)
|
||||
|
||||
headers = {
|
||||
'X-RateLimit-Limit': str(self.limit),
|
||||
'X-RateLimit-Remaining': str(remaining),
|
||||
'X-RateLimit-Reset': str(reset_time)
|
||||
}
|
||||
|
||||
return allowed, headers
|
||||
|
||||
|
||||
class SlidingWindow(RateLimitAlgorithm):
|
||||
"""Sliding window rate limiting algorithm."""
|
||||
|
||||
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
||||
key = self._get_key(identifier)
|
||||
now = time.time()
|
||||
|
||||
# Get current request timestamps
|
||||
timestamps = store.get(key) or []
|
||||
|
||||
# Remove timestamps outside the window
|
||||
window_start = now - self.window
|
||||
timestamps = [ts for ts in timestamps if ts > window_start]
|
||||
|
||||
# Check if request is allowed
|
||||
if len(timestamps) < self.limit:
|
||||
timestamps.append(now)
|
||||
allowed = True
|
||||
remaining = self.limit - len(timestamps)
|
||||
else:
|
||||
allowed = False
|
||||
remaining = 0
|
||||
|
||||
# Store updated timestamps
|
||||
store.set(key, timestamps, ttl=self.window + 1)
|
||||
|
||||
# Calculate reset time (when oldest request in window expires)
|
||||
if timestamps:
|
||||
reset_time = int(timestamps[0] + self.window)
|
||||
else:
|
||||
reset_time = int(now + self.window)
|
||||
|
||||
headers = {
|
||||
'X-RateLimit-Limit': str(self.limit),
|
||||
'X-RateLimit-Remaining': str(remaining),
|
||||
'X-RateLimit-Reset': str(reset_time)
|
||||
}
|
||||
|
||||
return allowed, headers
|
||||
|
||||
|
||||
class FixedWindow(RateLimitAlgorithm):
|
||||
"""Fixed window rate limiting algorithm."""
|
||||
|
||||
def is_allowed(self, identifier: str) -> Tuple[bool, Dict[str, Any]]:
|
||||
now = time.time()
|
||||
|
||||
# Calculate current window
|
||||
window_start = int(now // self.window) * self.window
|
||||
key = f"{self._get_key(identifier)}:{window_start}"
|
||||
|
||||
# Get current count for this window
|
||||
current_count = store.get(key) or 0
|
||||
|
||||
# Check if request is allowed
|
||||
if current_count < self.limit:
|
||||
new_count = store.increment(key, ttl=self.window + 1)
|
||||
allowed = True
|
||||
remaining = max(0, self.limit - new_count)
|
||||
else:
|
||||
allowed = False
|
||||
remaining = 0
|
||||
|
||||
# Calculate reset time (start of next window)
|
||||
reset_time = window_start + self.window
|
||||
|
||||
headers = {
|
||||
'X-RateLimit-Limit': str(self.limit),
|
||||
'X-RateLimit-Remaining': str(remaining),
|
||||
'X-RateLimit-Reset': str(int(reset_time))
|
||||
}
|
||||
|
||||
return allowed, headers
|
||||
@@ -0,0 +1,25 @@
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class LimitType(Enum):
|
||||
IP = "ip"
|
||||
API_KEY = "api_key"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimitConfig:
|
||||
"""Configuration for rate limiting."""
|
||||
limit: int
|
||||
window: int # Time window in seconds
|
||||
limit_type: LimitType = LimitType.IP
|
||||
api_key_header: str = "X-API-Key"
|
||||
error_message: str = "Rate limit exceeded"
|
||||
error_code: int = 429
|
||||
|
||||
def __post_init__(self):
|
||||
if self.limit <= 0:
|
||||
raise ValueError("Limit must be positive")
|
||||
if self.window <= 0:
|
||||
raise ValueError("Window must be positive")
|
||||
@@ -0,0 +1,105 @@
|
||||
from flask import Flask, jsonify
|
||||
from config import RateLimitConfig, LimitType
|
||||
from middleware import rate_limit
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# Rate limiting configurations
|
||||
token_bucket_config = RateLimitConfig(
|
||||
limit=10,
|
||||
window=60, # 10 requests per minute
|
||||
limit_type=LimitType.IP,
|
||||
error_message="Token bucket rate limit exceeded"
|
||||
)
|
||||
|
||||
sliding_window_config = RateLimitConfig(
|
||||
limit=5,
|
||||
window=30, # 5 requests per 30 seconds
|
||||
limit_type=LimitType.IP,
|
||||
error_message="Sliding window rate limit exceeded"
|
||||
)
|
||||
|
||||
fixed_window_config = RateLimitConfig(
|
||||
limit=20,
|
||||
window=60, # 20 requests per minute
|
||||
limit_type=LimitType.API_KEY,
|
||||
api_key_header="X-API-Key",
|
||||
error_message="Fixed window rate limit exceeded"
|
||||
)
|
||||
|
||||
|
||||
@app.route("/")
|
||||
def index():
|
||||
"""Home route without rate limiting."""
|
||||
return jsonify({
|
||||
"message": "Rate Limiting Example API",
|
||||
"endpoints": {
|
||||
"/token-bucket": "Token bucket algorithm (10 req/min by IP)",
|
||||
"/sliding-window": "Sliding window algorithm (5 req/30s by IP)",
|
||||
"/fixed-window": "Fixed window algorithm (20 req/min by API key)"
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
@app.route("/token-bucket")
|
||||
@rate_limit(token_bucket_config, algorithm="token_bucket")
|
||||
def token_bucket_endpoint():
|
||||
"""Endpoint protected by token bucket rate limiting."""
|
||||
return jsonify({
|
||||
"message": "Token bucket endpoint accessed successfully",
|
||||
"algorithm": "token_bucket",
|
||||
"limit": "10 requests per minute per IP"
|
||||
})
|
||||
|
||||
|
||||
@app.route("/sliding-window")
|
||||
@rate_limit(sliding_window_config, algorithm="sliding_window")
|
||||
def sliding_window_endpoint():
|
||||
"""Endpoint protected by sliding window rate limiting."""
|
||||
return jsonify({
|
||||
"message": "Sliding window endpoint accessed successfully",
|
||||
"algorithm": "sliding_window",
|
||||
"limit": "5 requests per 30 seconds per IP"
|
||||
})
|
||||
|
||||
|
||||
@app.route("/fixed-window")
|
||||
@rate_limit(fixed_window_config, algorithm="fixed_window")
|
||||
def fixed_window_endpoint():
|
||||
"""Endpoint protected by fixed window rate limiting."""
|
||||
return jsonify({
|
||||
"message": "Fixed window endpoint accessed successfully",
|
||||
"algorithm": "fixed_window",
|
||||
"limit": "20 requests per minute per API key"
|
||||
})
|
||||
|
||||
|
||||
@app.route("/health")
|
||||
def health():
|
||||
"""Health check endpoint without rate limiting."""
|
||||
return jsonify({"status": "healthy"})
|
||||
|
||||
|
||||
@app.errorhandler(429)
|
||||
def rate_limit_handler(e):
|
||||
"""Custom handler for rate limit errors."""
|
||||
return jsonify({
|
||||
"error": "Rate limit exceeded",
|
||||
"message": "Please wait before making more requests"
|
||||
}), 429
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Starting Flask app with rate limiting examples...")
|
||||
print("\nEndpoints:")
|
||||
print("- GET / : Home page with API information")
|
||||
print("- GET /token-bucket : Token bucket (10 req/min by IP)")
|
||||
print("- GET /sliding-window : Sliding window (5 req/30s by IP)")
|
||||
print("- GET /fixed-window : Fixed window (20 req/min by API key - requires X-API-Key header)")
|
||||
print("- GET /health : Health check (no rate limiting)")
|
||||
print("\nTesting with curl:")
|
||||
print("curl http://localhost:5000/token-bucket")
|
||||
print("curl -H 'X-API-Key: test-key' http://localhost:5000/fixed-window")
|
||||
print()
|
||||
|
||||
app.run(debug=True, host="0.0.0.0", port=5000)
|
||||
@@ -0,0 +1,94 @@
|
||||
import functools
|
||||
from typing import Union, Optional
|
||||
from flask import request, jsonify, Response
|
||||
from .algorithms import RateLimitAlgorithm, TokenBucket, SlidingWindow, FixedWindow
|
||||
from .config import RateLimitConfig, LimitType
|
||||
|
||||
|
||||
def rate_limit(
|
||||
config: RateLimitConfig,
|
||||
algorithm: str = "fixed_window"
|
||||
) -> callable:
|
||||
"""
|
||||
Flask decorator for rate limiting routes.
|
||||
|
||||
Args:
|
||||
config: RateLimitConfig instance
|
||||
algorithm: Algorithm type ("token_bucket", "sliding_window", "fixed_window")
|
||||
"""
|
||||
|
||||
# Create algorithm instance
|
||||
algorithm_map = {
|
||||
"token_bucket": TokenBucket,
|
||||
"sliding_window": SlidingWindow,
|
||||
"fixed_window": FixedWindow
|
||||
}
|
||||
|
||||
if algorithm not in algorithm_map:
|
||||
raise ValueError(f"Unknown algorithm: {algorithm}")
|
||||
|
||||
rate_limiter = algorithm_map[algorithm](config.limit, config.window)
|
||||
|
||||
def decorator(f):
|
||||
@functools.wraps(f)
|
||||
def wrapper(*args, **kwargs):
|
||||
# Get identifier based on limit type
|
||||
identifier = _get_identifier(config)
|
||||
|
||||
if identifier is None:
|
||||
return jsonify({
|
||||
"error": "Missing API key" if config.limit_type == LimitType.API_KEY else "Unable to identify client"
|
||||
}), 400
|
||||
|
||||
# Check rate limit
|
||||
allowed, headers = rate_limiter.is_allowed(identifier)
|
||||
|
||||
if not allowed:
|
||||
response = jsonify({"error": config.error_message})
|
||||
response.status_code = config.error_code
|
||||
|
||||
# Add rate limit headers
|
||||
for header_name, header_value in headers.items():
|
||||
response.headers[header_name] = header_value
|
||||
|
||||
return response
|
||||
|
||||
# Call original function
|
||||
result = f(*args, **kwargs)
|
||||
|
||||
# Add rate limit headers to successful response
|
||||
if isinstance(result, Response):
|
||||
response = result
|
||||
else:
|
||||
# Handle tuple returns (response, status_code, headers)
|
||||
if isinstance(result, tuple):
|
||||
if len(result) == 2:
|
||||
response = jsonify(result[0])
|
||||
response.status_code = result[1]
|
||||
elif len(result) == 3:
|
||||
response = jsonify(result[0])
|
||||
response.status_code = result[1]
|
||||
response.headers.update(result[2])
|
||||
else:
|
||||
response = jsonify(result)
|
||||
else:
|
||||
response = jsonify(result) if not isinstance(result, Response) else result
|
||||
|
||||
# Add rate limit headers
|
||||
for header_name, header_value in headers.items():
|
||||
response.headers[header_name] = header_value
|
||||
|
||||
return response
|
||||
|
||||
return wrapper
|
||||
return decorator
|
||||
|
||||
|
||||
def _get_identifier(config: RateLimitConfig) -> Optional[str]:
|
||||
"""Get client identifier based on configuration."""
|
||||
if config.limit_type == LimitType.IP:
|
||||
return request.environ.get('REMOTE_ADDR') or request.remote_addr
|
||||
elif config.limit_type == LimitType.API_KEY:
|
||||
return request.headers.get(config.api_key_header)
|
||||
else:
|
||||
return None
|
||||
@@ -0,0 +1 @@
|
||||
flask>=2.0.0
|
||||
@@ -0,0 +1,86 @@
|
||||
import time
|
||||
import threading
|
||||
from typing import Dict, Any, Optional
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
class InMemoryStore:
|
||||
"""Thread-safe in-memory store with TTL-based expiry."""
|
||||
|
||||
def __init__(self):
|
||||
self._data: Dict[str, Any] = {}
|
||||
self._expiry: Dict[str, float] = {}
|
||||
self._lock = threading.RLock()
|
||||
self._cleanup_interval = 60 # Cleanup every 60 seconds
|
||||
self._last_cleanup = time.time()
|
||||
|
||||
def get(self, key: str) -> Optional[Any]:
|
||||
"""Get value by key, returns None if expired or not found."""
|
||||
with self._lock:
|
||||
self._cleanup_if_needed()
|
||||
|
||||
if key not in self._data:
|
||||
return None
|
||||
|
||||
if key in self._expiry and time.time() > self._expiry[key]:
|
||||
self._delete_key(key)
|
||||
return None
|
||||
|
||||
return self._data[key]
|
||||
|
||||
def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
|
||||
"""Set value with optional TTL in seconds."""
|
||||
with self._lock:
|
||||
self._cleanup_if_needed()
|
||||
self._data[key] = value
|
||||
|
||||
if ttl is not None:
|
||||
self._expiry[key] = time.time() + ttl
|
||||
elif key in self._expiry:
|
||||
del self._expiry[key]
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
"""Delete key from store."""
|
||||
with self._lock:
|
||||
self._delete_key(key)
|
||||
|
||||
def increment(self, key: str, amount: int = 1, ttl: Optional[int] = None) -> int:
|
||||
"""Increment counter, creating if doesn't exist."""
|
||||
with self._lock:
|
||||
current = self.get(key) or 0
|
||||
new_value = current + amount
|
||||
self.set(key, new_value, ttl)
|
||||
return new_value
|
||||
|
||||
def _delete_key(self, key: str) -> None:
|
||||
"""Internal method to delete key without lock."""
|
||||
self._data.pop(key, None)
|
||||
self._expiry.pop(key, None)
|
||||
|
||||
def _cleanup_if_needed(self) -> None:
|
||||
"""Cleanup expired keys if interval has passed."""
|
||||
now = time.time()
|
||||
if now - self._last_cleanup > self._cleanup_interval:
|
||||
self._cleanup_expired()
|
||||
self._last_cleanup = now
|
||||
|
||||
def _cleanup_expired(self) -> None:
|
||||
"""Remove all expired keys."""
|
||||
now = time.time()
|
||||
expired_keys = [
|
||||
key for key, expiry_time in self._expiry.items()
|
||||
if now > expiry_time
|
||||
]
|
||||
|
||||
for key in expired_keys:
|
||||
self._delete_key(key)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all data."""
|
||||
with self._lock:
|
||||
self._data.clear()
|
||||
self._expiry.clear()
|
||||
|
||||
|
||||
# Global store instance
|
||||
store = InMemoryStore()
|
||||
@@ -0,0 +1,38 @@
|
||||
"""JSON Schema Validator Package
|
||||
|
||||
A from-scratch implementation of JSON schema validation supporting:
|
||||
- Basic JSON types (string, number, integer, boolean, array, object, null)
|
||||
- String constraints (minLength, maxLength, pattern)
|
||||
- Numeric constraints (minimum, maximum)
|
||||
- Enum constraints
|
||||
- Nested object and array validation
|
||||
- Required field validation
|
||||
- Detailed error reporting with JSON path tracking
|
||||
"""
|
||||
|
||||
from .validator import SchemaValidator
|
||||
from .errors import ValidationError
|
||||
from .types import (
|
||||
validate_string,
|
||||
validate_number,
|
||||
validate_integer,
|
||||
validate_boolean,
|
||||
validate_array,
|
||||
validate_object,
|
||||
validate_null
|
||||
)
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__author__ = "JSON Schema Validator"
|
||||
|
||||
__all__ = [
|
||||
"SchemaValidator",
|
||||
"ValidationError",
|
||||
"validate_string",
|
||||
"validate_number",
|
||||
"validate_integer",
|
||||
"validate_boolean",
|
||||
"validate_array",
|
||||
"validate_object",
|
||||
"validate_null"
|
||||
]
|
||||
@@ -0,0 +1,109 @@
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
from .validator import SchemaValidator
|
||||
from .errors import ValidationError
|
||||
|
||||
|
||||
def load_json_file(file_path: str) -> Dict[str, Any]:
|
||||
"""Load JSON data from file.
|
||||
|
||||
Args:
|
||||
file_path: Path to JSON file
|
||||
|
||||
Returns:
|
||||
Parsed JSON data
|
||||
|
||||
Raises:
|
||||
SystemExit: If file cannot be loaded or parsed
|
||||
"""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except FileNotFoundError:
|
||||
print(f"Error: File '{file_path}' not found.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Error: Invalid JSON in file '{file_path}': {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"Error: Failed to read file '{file_path}': {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def format_validation_results(errors: list, data_file: str) -> str:
|
||||
"""Format validation results for output.
|
||||
|
||||
Args:
|
||||
errors: List of ValidationError objects
|
||||
data_file: Name of the data file being validated
|
||||
|
||||
Returns:
|
||||
Formatted validation results
|
||||
"""
|
||||
if not errors:
|
||||
return f"✓ Validation successful: '{data_file}' is valid according to the schema."
|
||||
|
||||
result = f"✗ Validation failed: '{data_file}' has {len(errors)} error(s):\n\n"
|
||||
|
||||
for i, error in enumerate(errors, 1):
|
||||
result += f"{i}. {error}\n"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def main():
|
||||
"""Main CLI entry point."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="JSON Schema Validator - Validate JSON data against a schema"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--schema",
|
||||
required=True,
|
||||
help="Path to JSON schema file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data",
|
||||
required=True,
|
||||
help="Path to JSON data file to validate"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--format",
|
||||
choices=["human", "json"],
|
||||
default="human",
|
||||
help="Output format (default: human)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load schema and data files
|
||||
schema = load_json_file(args.schema)
|
||||
data = load_json_file(args.data)
|
||||
|
||||
# Create validator and perform validation
|
||||
try:
|
||||
validator = SchemaValidator(schema)
|
||||
errors = validator.validate(data)
|
||||
|
||||
# Output results based on format
|
||||
if args.format == "json":
|
||||
result = {
|
||||
"valid": len(errors) == 0,
|
||||
"errors": [error.to_dict() for error in errors]
|
||||
}
|
||||
print(json.dumps(result, indent=2))
|
||||
else:
|
||||
print(format_validation_results(errors, args.data))
|
||||
|
||||
# Exit with appropriate code
|
||||
sys.exit(0 if len(errors) == 0 else 1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error during validation: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class ValidationError(Exception):
|
||||
"""Custom exception for JSON schema validation errors."""
|
||||
|
||||
def __init__(self, path: str, message: str):
|
||||
"""Initialize validation error.
|
||||
|
||||
Args:
|
||||
path: JSON path to the failing field
|
||||
message: Human-readable error message
|
||||
"""
|
||||
self.path = path
|
||||
self.message = message
|
||||
super().__init__(f"Validation error at {path}: {message}")
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return formatted error message."""
|
||||
return f"Error at {self.path}: {self.message}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return detailed representation."""
|
||||
return f"ValidationError(path='{self.path}', message='{self.message}')"
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert error to dictionary representation.
|
||||
|
||||
Returns:
|
||||
Dictionary with path and message keys
|
||||
"""
|
||||
return {
|
||||
"path": self.path,
|
||||
"message": self.message
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
# No external dependencies required
|
||||
# This JSON schema validator uses only Python standard library
|
||||
@@ -0,0 +1,208 @@
|
||||
import re
|
||||
from typing import Any, Dict
|
||||
from .errors import ValidationError
|
||||
|
||||
|
||||
def validate_string(value: Any, schema: Dict[str, Any], path: str) -> None:
|
||||
"""Validate string type and constraints.
|
||||
|
||||
Args:
|
||||
value: Value to validate
|
||||
schema: Schema definition
|
||||
path: JSON path to the value
|
||||
|
||||
Raises:
|
||||
ValidationError: If validation fails
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Expected string, got {type(value).__name__}"
|
||||
)
|
||||
|
||||
# Check minLength constraint
|
||||
min_length = schema.get("minLength")
|
||||
if min_length is not None and len(value) < min_length:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"String length {len(value)} is less than minimum {min_length}"
|
||||
)
|
||||
|
||||
# Check maxLength constraint
|
||||
max_length = schema.get("maxLength")
|
||||
if max_length is not None and len(value) > max_length:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"String length {len(value)} exceeds maximum {max_length}"
|
||||
)
|
||||
|
||||
# Check pattern constraint
|
||||
pattern = schema.get("pattern")
|
||||
if pattern is not None:
|
||||
try:
|
||||
if not re.match(pattern, value):
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"String '{value}' does not match pattern '{pattern}'"
|
||||
)
|
||||
except re.error as e:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Invalid regex pattern '{pattern}': {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
def validate_number(value: Any, schema: Dict[str, Any], path: str) -> None:
|
||||
"""Validate number type and constraints.
|
||||
|
||||
Args:
|
||||
value: Value to validate
|
||||
schema: Schema definition
|
||||
path: JSON path to the value
|
||||
|
||||
Raises:
|
||||
ValidationError: If validation fails
|
||||
"""
|
||||
if not isinstance(value, (int, float)) or isinstance(value, bool):
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Expected number, got {type(value).__name__}"
|
||||
)
|
||||
|
||||
# Check minimum constraint
|
||||
minimum = schema.get("minimum")
|
||||
if minimum is not None and value < minimum:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Value {value} is less than minimum {minimum}"
|
||||
)
|
||||
|
||||
# Check maximum constraint
|
||||
maximum = schema.get("maximum")
|
||||
if maximum is not None and value > maximum:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Value {value} exceeds maximum {maximum}"
|
||||
)
|
||||
|
||||
|
||||
def validate_integer(value: Any, schema: Dict[str, Any], path: str) -> None:
|
||||
"""Validate integer type and constraints.
|
||||
|
||||
Args:
|
||||
value: Value to validate
|
||||
schema: Schema definition
|
||||
path: JSON path to the value
|
||||
|
||||
Raises:
|
||||
ValidationError: If validation fails
|
||||
"""
|
||||
if not isinstance(value, int) or isinstance(value, bool):
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Expected integer, got {type(value).__name__}"
|
||||
)
|
||||
|
||||
# Check minimum constraint
|
||||
minimum = schema.get("minimum")
|
||||
if minimum is not None and value < minimum:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Value {value} is less than minimum {minimum}"
|
||||
)
|
||||
|
||||
# Check maximum constraint
|
||||
maximum = schema.get("maximum")
|
||||
if maximum is not None and value > maximum:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Value {value} exceeds maximum {maximum}"
|
||||
)
|
||||
|
||||
|
||||
def validate_boolean(value: Any, schema: Dict[str, Any], path: str) -> None:
|
||||
"""Validate boolean type.
|
||||
|
||||
Args:
|
||||
value: Value to validate
|
||||
schema: Schema definition
|
||||
path: JSON path to the value
|
||||
|
||||
Raises:
|
||||
ValidationError: If validation fails
|
||||
"""
|
||||
if not isinstance(value, bool):
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Expected boolean, got {type(value).__name__}"
|
||||
)
|
||||
|
||||
|
||||
def validate_array(value: Any, schema: Dict[str, Any], path: str) -> None:
|
||||
"""Validate array type and constraints.
|
||||
|
||||
Args:
|
||||
value: Value to validate
|
||||
schema: Schema definition
|
||||
path: JSON path to the value
|
||||
|
||||
Raises:
|
||||
ValidationError: If validation fails
|
||||
"""
|
||||
if not isinstance(value, list):
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Expected array, got {type(value).__name__}"
|
||||
)
|
||||
|
||||
# Check minLength constraint for arrays
|
||||
min_length = schema.get("minLength")
|
||||
if min_length is not None and len(value) < min_length:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Array length {len(value)} is less than minimum {min_length}"
|
||||
)
|
||||
|
||||
# Check maxLength constraint for arrays
|
||||
max_length = schema.get("maxLength")
|
||||
if max_length is not None and len(value) > max_length:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Array length {len(value)} exceeds maximum {max_length}"
|
||||
)
|
||||
|
||||
|
||||
def validate_object(value: Any, schema: Dict[str, Any], path: str) -> None:
|
||||
"""Validate object type.
|
||||
|
||||
Args:
|
||||
value: Value to validate
|
||||
schema: Schema definition
|
||||
path: JSON path to the value
|
||||
|
||||
Raises:
|
||||
ValidationError: If validation fails
|
||||
"""
|
||||
if not isinstance(value, dict):
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Expected object, got {type(value).__name__}"
|
||||
)
|
||||
|
||||
|
||||
def validate_null(value: Any, schema: Dict[str, Any], path: str) -> None:
|
||||
"""Validate null type.
|
||||
|
||||
Args:
|
||||
value: Value to validate
|
||||
schema: Schema definition
|
||||
path: JSON path to the value
|
||||
|
||||
Raises:
|
||||
ValidationError: If validation fails
|
||||
"""
|
||||
if value is not None:
|
||||
raise ValidationError(
|
||||
path=path,
|
||||
message=f"Expected null, got {type(value).__name__}"
|
||||
)
|
||||
@@ -0,0 +1,152 @@
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from .errors import ValidationError
|
||||
from .types import validate_string, validate_number, validate_integer, validate_boolean, validate_array, validate_object, validate_null
|
||||
|
||||
|
||||
class SchemaValidator:
|
||||
"""Main JSON schema validator class."""
|
||||
|
||||
def __init__(self, schema: Dict[str, Any]):
|
||||
"""Initialize validator with a JSON schema.
|
||||
|
||||
Args:
|
||||
schema: JSON schema dictionary
|
||||
"""
|
||||
self.schema = schema
|
||||
|
||||
def validate(self, data: Any) -> List[ValidationError]:
|
||||
"""Validate data against the schema.
|
||||
|
||||
Args:
|
||||
data: Data to validate
|
||||
|
||||
Returns:
|
||||
List of validation errors (empty if valid)
|
||||
"""
|
||||
errors = []
|
||||
self._validate_recursive(data, self.schema, "$", errors)
|
||||
return errors
|
||||
|
||||
def _validate_recursive(self, data: Any, schema: Dict[str, Any], path: str, errors: List[ValidationError]) -> None:
|
||||
"""Recursively validate data against schema.
|
||||
|
||||
Args:
|
||||
data: Current data being validated
|
||||
schema: Current schema definition
|
||||
path: JSON path to current data
|
||||
errors: List to accumulate errors
|
||||
"""
|
||||
# Handle type validation
|
||||
schema_type = schema.get("type")
|
||||
if schema_type:
|
||||
if schema_type == "string":
|
||||
self._validate_type(validate_string, data, schema, path, errors)
|
||||
elif schema_type == "number":
|
||||
self._validate_type(validate_number, data, schema, path, errors)
|
||||
elif schema_type == "integer":
|
||||
self._validate_type(validate_integer, data, schema, path, errors)
|
||||
elif schema_type == "boolean":
|
||||
self._validate_type(validate_boolean, data, schema, path, errors)
|
||||
elif schema_type == "array":
|
||||
self._validate_array(data, schema, path, errors)
|
||||
elif schema_type == "object":
|
||||
self._validate_object(data, schema, path, errors)
|
||||
elif schema_type == "null":
|
||||
self._validate_type(validate_null, data, schema, path, errors)
|
||||
else:
|
||||
errors.append(ValidationError(
|
||||
path=path,
|
||||
message=f"Unknown type '{schema_type}' in schema"
|
||||
))
|
||||
|
||||
# Handle enum constraint (applies to all types)
|
||||
if "enum" in schema:
|
||||
if data not in schema["enum"]:
|
||||
errors.append(ValidationError(
|
||||
path=path,
|
||||
message=f"Value must be one of {schema['enum']}, got {data}"
|
||||
))
|
||||
|
||||
def _validate_type(self, validator_func, data: Any, schema: Dict[str, Any], path: str, errors: List[ValidationError]) -> None:
|
||||
"""Validate data using a type-specific validator function.
|
||||
|
||||
Args:
|
||||
validator_func: Type-specific validation function
|
||||
data: Data to validate
|
||||
schema: Schema definition
|
||||
path: JSON path
|
||||
errors: List to accumulate errors
|
||||
"""
|
||||
try:
|
||||
validator_func(data, schema, path)
|
||||
except ValidationError as e:
|
||||
errors.append(e)
|
||||
|
||||
def _validate_array(self, data: Any, schema: Dict[str, Any], path: str, errors: List[ValidationError]) -> None:
|
||||
"""Validate array data and its items.
|
||||
|
||||
Args:
|
||||
data: Data to validate
|
||||
schema: Schema definition
|
||||
path: JSON path
|
||||
errors: List to accumulate errors
|
||||
"""
|
||||
try:
|
||||
validate_array(data, schema, path)
|
||||
except ValidationError as e:
|
||||
errors.append(e)
|
||||
return
|
||||
|
||||
# Validate array items if schema is provided
|
||||
items_schema = schema.get("items")
|
||||
if items_schema and isinstance(data, list):
|
||||
for i, item in enumerate(data):
|
||||
item_path = f"{path}[{i}]"
|
||||
self._validate_recursive(item, items_schema, item_path, errors)
|
||||
|
||||
def _validate_object(self, data: Any, schema: Dict[str, Any], path: str, errors: List[ValidationError]) -> None:
|
||||
"""Validate object data and its properties.
|
||||
|
||||
Args:
|
||||
data: Data to validate
|
||||
schema: Schema definition
|
||||
path: JSON path
|
||||
errors: List to accumulate errors
|
||||
"""
|
||||
try:
|
||||
validate_object(data, schema, path)
|
||||
except ValidationError as e:
|
||||
errors.append(e)
|
||||
return
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return
|
||||
|
||||
# Check required fields
|
||||
required = schema.get("required", [])
|
||||
for field in required:
|
||||
if field not in data:
|
||||
field_path = f"{path}.{field}" if path != "$" else f"$.{field}"
|
||||
errors.append(ValidationError(
|
||||
path=field_path,
|
||||
message=f"Required field '{field}' is missing"
|
||||
))
|
||||
|
||||
# Validate object properties
|
||||
properties = schema.get("properties", {})
|
||||
for field, value in data.items():
|
||||
if field in properties:
|
||||
field_path = f"{path}.{field}" if path != "$" else f"$.{field}"
|
||||
self._validate_recursive(value, properties[field], field_path, errors)
|
||||
|
||||
def is_valid(self, data: Any) -> bool:
|
||||
"""Check if data is valid against the schema.
|
||||
|
||||
Args:
|
||||
data: Data to validate
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
"""
|
||||
return len(self.validate(data)) == 0
|
||||
@@ -0,0 +1,185 @@
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from jsonpath_ng import parse as jsonpath_parse
|
||||
|
||||
|
||||
class AssertionResult:
|
||||
"""Result of a single assertion."""
|
||||
|
||||
def __init__(self, passed: bool, message: str, assertion_type: str):
|
||||
self.passed = passed
|
||||
self.message = message
|
||||
self.assertion_type = assertion_type
|
||||
|
||||
|
||||
class ResponseValidator:
|
||||
"""Validates HTTP responses against expected criteria."""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def validate_response(self, response, expected: Dict[str, Any]) -> List[AssertionResult]:
|
||||
"""Validate response against all expected criteria.
|
||||
|
||||
Args:
|
||||
response: HTTP response object
|
||||
expected: Dictionary containing expected response criteria
|
||||
|
||||
Returns:
|
||||
List of AssertionResult objects
|
||||
"""
|
||||
results = []
|
||||
|
||||
# Validate status code
|
||||
if "status" in expected:
|
||||
results.append(self._validate_status_code(response, expected["status"]))
|
||||
|
||||
# Validate body contains
|
||||
if "body_contains" in expected:
|
||||
body_contains = expected["body_contains"]
|
||||
if isinstance(body_contains, str):
|
||||
body_contains = [body_contains]
|
||||
for text in body_contains:
|
||||
results.append(self._validate_body_contains(response, text))
|
||||
|
||||
# Validate JSON path assertions
|
||||
if "json_path" in expected:
|
||||
json_assertions = expected["json_path"]
|
||||
if isinstance(json_assertions, dict):
|
||||
json_assertions = [json_assertions]
|
||||
for assertion in json_assertions:
|
||||
results.append(self._validate_json_path(response, assertion))
|
||||
|
||||
# Validate headers
|
||||
if "headers" in expected:
|
||||
for header_name, expected_value in expected["headers"].items():
|
||||
results.append(self._validate_header(response, header_name, expected_value))
|
||||
|
||||
# Validate response time (if specified)
|
||||
if "max_response_time" in expected:
|
||||
results.append(self._validate_response_time(response, expected["max_response_time"]))
|
||||
|
||||
return results
|
||||
|
||||
def _validate_status_code(self, response, expected_status: Union[int, List[int]]) -> AssertionResult:
|
||||
"""Validate HTTP status code."""
|
||||
if isinstance(expected_status, list):
|
||||
passed = response.status_code in expected_status
|
||||
expected_str = f"one of {expected_status}"
|
||||
else:
|
||||
passed = response.status_code == expected_status
|
||||
expected_str = str(expected_status)
|
||||
|
||||
if passed:
|
||||
message = f"Status code {response.status_code} matches expected {expected_str}"
|
||||
else:
|
||||
message = f"Expected status {expected_str}, got {response.status_code}"
|
||||
|
||||
return AssertionResult(passed, message, "status_code")
|
||||
|
||||
def _validate_body_contains(self, response, expected_text: str) -> AssertionResult:
|
||||
"""Validate that response body contains expected text."""
|
||||
response_text = response.text
|
||||
passed = expected_text in response_text
|
||||
|
||||
if passed:
|
||||
message = f"Response body contains '{expected_text}'"
|
||||
else:
|
||||
message = f"Response body does not contain '{expected_text}'"
|
||||
|
||||
return AssertionResult(passed, message, "body_contains")
|
||||
|
||||
def _validate_json_path(self, response, assertion: Dict[str, Any]) -> AssertionResult:
|
||||
"""Validate JSON path assertion."""
|
||||
path = assertion.get("path", "")
|
||||
expected_value = assertion.get("value")
|
||||
operator = assertion.get("operator", "equals")
|
||||
|
||||
try:
|
||||
json_data = response.json()
|
||||
jsonpath_expr = jsonpath_parse(path)
|
||||
matches = jsonpath_expr.find(json_data)
|
||||
|
||||
if not matches:
|
||||
return AssertionResult(
|
||||
False,
|
||||
f"JSON path '{path}' not found in response",
|
||||
"json_path"
|
||||
)
|
||||
|
||||
actual_value = matches[0].value
|
||||
passed = self._compare_values(actual_value, expected_value, operator)
|
||||
|
||||
if passed:
|
||||
message = f"JSON path '{path}' assertion passed: {actual_value} {operator} {expected_value}"
|
||||
else:
|
||||
message = f"JSON path '{path}' assertion failed: {actual_value} {operator} {expected_value}"
|
||||
|
||||
return AssertionResult(passed, message, "json_path")
|
||||
|
||||
except json.JSONDecodeError:
|
||||
return AssertionResult(
|
||||
False,
|
||||
f"Response is not valid JSON for path '{path}'",
|
||||
"json_path"
|
||||
)
|
||||
except Exception as e:
|
||||
return AssertionResult(
|
||||
False,
|
||||
f"JSON path assertion error: {str(e)}",
|
||||
"json_path"
|
||||
)
|
||||
|
||||
def _validate_header(self, response, header_name: str, expected_value: str) -> AssertionResult:
|
||||
"""Validate response header value."""
|
||||
actual_value = response.headers.get(header_name)
|
||||
|
||||
if actual_value is None:
|
||||
return AssertionResult(
|
||||
False,
|
||||
f"Header '{header_name}' not found in response",
|
||||
"header"
|
||||
)
|
||||
|
||||
passed = str(actual_value) == str(expected_value)
|
||||
|
||||
if passed:
|
||||
message = f"Header '{header_name}' matches expected value '{expected_value}'"
|
||||
else:
|
||||
message = f"Header '{header_name}' expected '{expected_value}', got '{actual_value}'"
|
||||
|
||||
return AssertionResult(passed, message, "header")
|
||||
|
||||
def _validate_response_time(self, response, max_time: float) -> AssertionResult:
|
||||
"""Validate response time is within acceptable limit."""
|
||||
response_time = response.elapsed.total_seconds()
|
||||
passed = response_time <= max_time
|
||||
|
||||
if passed:
|
||||
message = f"Response time {response_time:.3f}s is within limit of {max_time}s"
|
||||
else:
|
||||
message = f"Response time {response_time:.3f}s exceeds limit of {max_time}s"
|
||||
|
||||
return AssertionResult(passed, message, "response_time")
|
||||
|
||||
def _compare_values(self, actual: Any, expected: Any, operator: str) -> bool:
|
||||
"""Compare two values using specified operator."""
|
||||
if operator == "equals":
|
||||
return actual == expected
|
||||
elif operator == "not_equals":
|
||||
return actual != expected
|
||||
elif operator == "greater_than":
|
||||
return float(actual) > float(expected)
|
||||
elif operator == "less_than":
|
||||
return float(actual) < float(expected)
|
||||
elif operator == "greater_than_or_equal":
|
||||
return float(actual) >= float(expected)
|
||||
elif operator == "less_than_or_equal":
|
||||
return float(actual) <= float(expected)
|
||||
elif operator == "contains":
|
||||
return str(expected) in str(actual)
|
||||
elif operator == "regex":
|
||||
return bool(re.search(str(expected), str(actual)))
|
||||
else:
|
||||
raise ValueError(f"Unknown operator: {operator}")
|
||||
@@ -0,0 +1,98 @@
|
||||
import argparse
|
||||
import sys
|
||||
from runner import TestRunner
|
||||
|
||||
|
||||
def main():
|
||||
"""Command-line interface for HTTP API testing framework."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="HTTP API Testing Framework",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python cli.py --suite tests.yaml
|
||||
python cli.py --suite tests.yaml --base-url https://api.example.com
|
||||
python cli.py --suite tests.yaml --base-url https://api.example.com --verbose
|
||||
python cli.py --suite tests.yaml --output report.txt
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--suite',
|
||||
required=True,
|
||||
help='Path to YAML test suite file'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--base-url',
|
||||
help='Base URL for API endpoints (can be overridden by suite file)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--verbose', '-v',
|
||||
action='store_true',
|
||||
help='Show detailed test results including passed tests'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--output', '-o',
|
||||
help='Save report to specified file'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--quiet', '-q',
|
||||
action='store_true',
|
||||
help='Suppress console output except for final summary'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
# Initialize test runner
|
||||
runner = TestRunner(base_url=args.base_url or "")
|
||||
|
||||
if not args.quiet:
|
||||
print(f"Loading test suite: {args.suite}")
|
||||
if args.base_url:
|
||||
print(f"Base URL: {args.base_url}")
|
||||
print("Starting test execution...")
|
||||
print()
|
||||
|
||||
# Execute tests
|
||||
reporter = runner.run_test_suite(args.suite)
|
||||
|
||||
# Generate and display report
|
||||
if not args.quiet:
|
||||
print(reporter.generate_report(verbose=args.verbose))
|
||||
else:
|
||||
reporter.print_summary()
|
||||
|
||||
# Save report to file if requested
|
||||
if args.output:
|
||||
reporter.save_report(args.output, verbose=args.verbose)
|
||||
if not args.quiet:
|
||||
print(f"Report saved to: {args.output}")
|
||||
|
||||
# Exit with appropriate code
|
||||
failed_tests = len(reporter.get_failed_tests())
|
||||
if failed_tests > 0:
|
||||
sys.exit(1)
|
||||
else:
|
||||
sys.exit(0)
|
||||
|
||||
except FileNotFoundError as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
sys.exit(2)
|
||||
except ValueError as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
sys.exit(2)
|
||||
except KeyboardInterrupt:
|
||||
print("\nTest execution interrupted by user", file=sys.stderr)
|
||||
sys.exit(130)
|
||||
except Exception as e:
|
||||
print(f"Unexpected error: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,137 @@
|
||||
import time
|
||||
from typing import List, Dict, Any
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestResult:
|
||||
"""Result of a single test case execution."""
|
||||
name: str
|
||||
method: str
|
||||
url: str
|
||||
status_code: int
|
||||
response_time: float
|
||||
passed: bool
|
||||
assertion_results: List[Any] # List of AssertionResult objects
|
||||
error: str = None
|
||||
|
||||
|
||||
class TestReporter:
|
||||
"""Generates test execution reports with detailed pass/fail information."""
|
||||
|
||||
def __init__(self):
|
||||
self.results: List[TestResult] = []
|
||||
self.start_time = None
|
||||
self.end_time = None
|
||||
|
||||
def start_execution(self):
|
||||
"""Mark the start of test execution."""
|
||||
self.start_time = time.time()
|
||||
|
||||
def end_execution(self):
|
||||
"""Mark the end of test execution."""
|
||||
self.end_time = time.time()
|
||||
|
||||
def add_result(self, result: TestResult):
|
||||
"""Add a test result to the report."""
|
||||
self.results.append(result)
|
||||
|
||||
def generate_report(self, verbose: bool = True) -> str:
|
||||
"""Generate formatted test report.
|
||||
|
||||
Args:
|
||||
verbose: Whether to include detailed failure information
|
||||
|
||||
Returns:
|
||||
Formatted report string
|
||||
"""
|
||||
report_lines = []
|
||||
|
||||
# Header
|
||||
report_lines.append("=" * 80)
|
||||
report_lines.append("HTTP API TEST RESULTS")
|
||||
report_lines.append("=" * 80)
|
||||
report_lines.append("")
|
||||
|
||||
# Summary statistics
|
||||
total_tests = len(self.results)
|
||||
passed_tests = sum(1 for r in self.results if r.passed)
|
||||
failed_tests = total_tests - passed_tests
|
||||
|
||||
execution_time = (self.end_time - self.start_time) if (self.end_time and self.start_time) else 0
|
||||
|
||||
report_lines.append(f"Total Tests: {total_tests}")
|
||||
report_lines.append(f"Passed: {passed_tests}")
|
||||
report_lines.append(f"Failed: {failed_tests}")
|
||||
report_lines.append(f"Success Rate: {(passed_tests/total_tests*100):.1f}%" if total_tests > 0 else "Success Rate: 0%")
|
||||
report_lines.append(f"Execution Time: {execution_time:.2f}s")
|
||||
report_lines.append("")
|
||||
|
||||
# Test results
|
||||
if verbose:
|
||||
for i, result in enumerate(self.results, 1):
|
||||
status_icon = "✓" if result.passed else "✗"
|
||||
report_lines.append(f"{i}. [{status_icon}] {result.name}")
|
||||
report_lines.append(f" {result.method} {result.url}")
|
||||
report_lines.append(f" Status: {result.status_code} | Response Time: {result.response_time:.3f}s")
|
||||
|
||||
if result.error:
|
||||
report_lines.append(f" ERROR: {result.error}")
|
||||
|
||||
# Show assertion details for failed tests
|
||||
if not result.passed and result.assertion_results:
|
||||
report_lines.append(" Assertion Results:")
|
||||
for assertion in result.assertion_results:
|
||||
assertion_icon = "✓" if assertion.passed else "✗"
|
||||
report_lines.append(f" [{assertion_icon}] {assertion.message}")
|
||||
|
||||
report_lines.append("")
|
||||
else:
|
||||
# Compact view - only show failed tests
|
||||
failed_results = [r for r in self.results if not r.passed]
|
||||
if failed_results:
|
||||
report_lines.append("FAILED TESTS:")
|
||||
report_lines.append("-" * 40)
|
||||
for result in failed_results:
|
||||
report_lines.append(f"✗ {result.name}")
|
||||
report_lines.append(f" {result.method} {result.url} -> {result.status_code}")
|
||||
if result.error:
|
||||
report_lines.append(f" ERROR: {result.error}")
|
||||
report_lines.append("")
|
||||
|
||||
# Footer
|
||||
report_lines.append("=" * 80)
|
||||
|
||||
return "\n".join(report_lines)
|
||||
|
||||
def print_summary(self):
|
||||
"""Print a quick summary to console."""
|
||||
total = len(self.results)
|
||||
passed = sum(1 for r in self.results if r.passed)
|
||||
failed = total - passed
|
||||
|
||||
print(f"\nTest Summary: {passed}/{total} passed, {failed} failed")
|
||||
if failed > 0:
|
||||
print("Failed tests:")
|
||||
for result in self.results:
|
||||
if not result.passed:
|
||||
print(f" - {result.name}: {result.error or 'Assertion failures'}")
|
||||
|
||||
def get_failed_tests(self) -> List[TestResult]:
|
||||
"""Get list of failed test results."""
|
||||
return [r for r in self.results if not r.passed]
|
||||
|
||||
def get_passed_tests(self) -> List[TestResult]:
|
||||
"""Get list of passed test results."""
|
||||
return [r for r in self.results if r.passed]
|
||||
|
||||
def get_success_rate(self) -> float:
|
||||
"""Get success rate as percentage."""
|
||||
if not self.results:
|
||||
return 0.0
|
||||
return (len(self.get_passed_tests()) / len(self.results)) * 100
|
||||
|
||||
def save_report(self, filename: str, verbose: bool = True):
|
||||
"""Save report to file."""
|
||||
with open(filename, 'w') as f:
|
||||
f.write(self.generate_report(verbose))
|
||||
@@ -0,0 +1,3 @@
|
||||
requests>=2.28.0
|
||||
pyyaml>=6.0
|
||||
jsonpath-ng>=1.5.0
|
||||
@@ -0,0 +1,186 @@
|
||||
import yaml
|
||||
import requests
|
||||
import json
|
||||
import time
|
||||
from typing import Dict, List, Any, Optional
|
||||
from urllib.parse import urljoin
|
||||
|
||||
from variables import VariableContext
|
||||
from assertions import ResponseValidator
|
||||
from reporter import TestReporter, TestResult
|
||||
|
||||
|
||||
class TestRunner:
|
||||
"""Main test execution engine for HTTP API testing framework."""
|
||||
|
||||
def __init__(self, base_url: str = ""):
|
||||
self.base_url = base_url.rstrip('/')
|
||||
self.variable_context = VariableContext()
|
||||
self.validator = ResponseValidator()
|
||||
self.reporter = TestReporter()
|
||||
self.session = requests.Session()
|
||||
|
||||
def load_test_suite(self, suite_path: str) -> Dict[str, Any]:
|
||||
"""Load test suite from YAML file.
|
||||
|
||||
Args:
|
||||
suite_path: Path to YAML test suite file
|
||||
|
||||
Returns:
|
||||
Parsed test suite dictionary
|
||||
"""
|
||||
try:
|
||||
with open(suite_path, 'r', encoding='utf-8') as f:
|
||||
suite = yaml.safe_load(f)
|
||||
return suite
|
||||
except FileNotFoundError:
|
||||
raise FileNotFoundError(f"Test suite file not found: {suite_path}")
|
||||
except yaml.YAMLError as e:
|
||||
raise ValueError(f"Invalid YAML in test suite: {e}")
|
||||
|
||||
def run_test_suite(self, suite_path: str) -> TestReporter:
|
||||
"""Execute complete test suite.
|
||||
|
||||
Args:
|
||||
suite_path: Path to YAML test suite file
|
||||
|
||||
Returns:
|
||||
TestReporter with execution results
|
||||
"""
|
||||
suite = self.load_test_suite(suite_path)
|
||||
|
||||
# Set up suite-level variables
|
||||
suite_vars = suite.get('variables', {})
|
||||
for name, value in suite_vars.items():
|
||||
self.variable_context.set_variable(name, value)
|
||||
|
||||
# Update base URL if specified in suite
|
||||
if 'base_url' in suite and not self.base_url:
|
||||
self.base_url = suite['base_url'].rstrip('/')
|
||||
|
||||
# Execute tests
|
||||
self.reporter.start_execution()
|
||||
|
||||
tests = suite.get('tests', [])
|
||||
for test_case in tests:
|
||||
try:
|
||||
result = self.run_single_test(test_case)
|
||||
self.reporter.add_result(result)
|
||||
except Exception as e:
|
||||
# Create failed result for test execution errors
|
||||
error_result = TestResult(
|
||||
name=test_case.get('name', 'Unknown Test'),
|
||||
method=test_case.get('method', 'GET'),
|
||||
url=test_case.get('url', ''),
|
||||
status_code=0,
|
||||
response_time=0.0,
|
||||
passed=False,
|
||||
assertion_results=[],
|
||||
error=str(e)
|
||||
)
|
||||
self.reporter.add_result(error_result)
|
||||
|
||||
self.reporter.end_execution()
|
||||
return self.reporter
|
||||
|
||||
def run_single_test(self, test_case: Dict[str, Any]) -> TestResult:
|
||||
"""Execute a single test case.
|
||||
|
||||
Args:
|
||||
test_case: Test case configuration dictionary
|
||||
|
||||
Returns:
|
||||
TestResult object
|
||||
"""
|
||||
test_name = test_case.get('name', 'Unnamed Test')
|
||||
|
||||
# Build request parameters with variable substitution
|
||||
method = test_case.get('method', 'GET').upper()
|
||||
url = self._build_url(test_case.get('url', ''))
|
||||
headers = self._substitute_variables_in_data(test_case.get('headers', {}))
|
||||
|
||||
# Handle request body
|
||||
body = test_case.get('body')
|
||||
json_body = test_case.get('json')
|
||||
|
||||
request_kwargs = {
|
||||
'headers': headers,
|
||||
'timeout': test_case.get('timeout', 30)
|
||||
}
|
||||
|
||||
if json_body is not None:
|
||||
request_kwargs['json'] = self._substitute_variables_in_data(json_body)
|
||||
elif body is not None:
|
||||
if isinstance(body, str):
|
||||
request_kwargs['data'] = self.variable_context.substitute_variables(body)
|
||||
else:
|
||||
request_kwargs['data'] = self._substitute_variables_in_data(body)
|
||||
|
||||
# Add query parameters if specified
|
||||
params = test_case.get('params', {})
|
||||
if params:
|
||||
request_kwargs['params'] = self._substitute_variables_in_data(params)
|
||||
|
||||
# Execute request
|
||||
start_time = time.time()
|
||||
try:
|
||||
response = self.session.request(method, url, **request_kwargs)
|
||||
response_time = time.time() - start_time
|
||||
except requests.RequestException as e:
|
||||
return TestResult(
|
||||
name=test_name,
|
||||
method=method,
|
||||
url=url,
|
||||
status_code=0,
|
||||
response_time=time.time() - start_time,
|
||||
passed=False,
|
||||
assertion_results=[],
|
||||
error=f"Request failed: {str(e)}"
|
||||
)
|
||||
|
||||
# Extract variables from response
|
||||
extract_config = test_case.get('extract', {})
|
||||
if extract_config:
|
||||
self.variable_context.extract_variables(response, extract_config)
|
||||
|
||||
# Validate response
|
||||
expected = test_case.get('expected', {})
|
||||
assertion_results = []
|
||||
if expected:
|
||||
assertion_results = self.validator.validate_response(response, expected)
|
||||
|
||||
# Determine overall test result
|
||||
passed = all(assertion.passed for assertion in assertion_results)
|
||||
|
||||
return TestResult(
|
||||
name=test_name,
|
||||
method=method,
|
||||
url=url,
|
||||
status_code=response.status_code,
|
||||
response_time=response_time,
|
||||
passed=passed,
|
||||
assertion_results=assertion_results
|
||||
)
|
||||
|
||||
def _build_url(self, path: str) -> str:
|
||||
"""Build complete URL from base URL and path."""
|
||||
path = self.variable_context.substitute_variables(path)
|
||||
|
||||
if path.startswith(('http://', 'https://')):
|
||||
return path
|
||||
|
||||
if self.base_url:
|
||||
return urljoin(self.base_url + '/', path.lstrip('/'))
|
||||
|
||||
return path
|
||||
|
||||
def _substitute_variables_in_data(self, data: Any) -> Any:
|
||||
"""Substitute variables in request data."""
|
||||
if isinstance(data, dict):
|
||||
return self.variable_context.substitute_in_dict(data)
|
||||
elif isinstance(data, list):
|
||||
return self.variable_context.substitute_in_list(data)
|
||||
elif isinstance(data, str):
|
||||
return self.variable_context.substitute_variables(data)
|
||||
else:
|
||||
return data
|
||||
@@ -0,0 +1,153 @@
|
||||
# Example HTTP API Test Suite
|
||||
# This file demonstrates the YAML format for defining test cases
|
||||
|
||||
# Suite-level configuration
|
||||
base_url: "https://jsonplaceholder.typicode.com"
|
||||
|
||||
# Global variables available to all tests
|
||||
variables:
|
||||
api_version: "v1"
|
||||
default_timeout: 30
|
||||
|
||||
# Test cases executed sequentially
|
||||
tests:
|
||||
- name: "Get all posts"
|
||||
method: GET
|
||||
url: "/posts"
|
||||
expected:
|
||||
status: 200
|
||||
body_contains: "userId"
|
||||
json_path:
|
||||
- path: "$[0].id"
|
||||
value: 1
|
||||
operator: "equals"
|
||||
- path: "$"
|
||||
operator: "greater_than"
|
||||
value: 0
|
||||
extract:
|
||||
first_post_id:
|
||||
type: "json_path"
|
||||
path: "$[0].id"
|
||||
user_id:
|
||||
type: "json_path"
|
||||
path: "$[0].userId"
|
||||
|
||||
- name: "Get specific post using extracted ID"
|
||||
method: GET
|
||||
url: "/posts/${first_post_id}"
|
||||
expected:
|
||||
status: 200
|
||||
json_path:
|
||||
- path: "$.id"
|
||||
value: "${first_post_id}"
|
||||
operator: "equals"
|
||||
- path: "$.userId"
|
||||
value: "${user_id}"
|
||||
operator: "equals"
|
||||
extract:
|
||||
post_title:
|
||||
type: "json_path"
|
||||
path: "$.title"
|
||||
|
||||
- name: "Create new post"
|
||||
method: POST
|
||||
url: "/posts"
|
||||
headers:
|
||||
Content-Type: "application/json"
|
||||
json:
|
||||
title: "Test Post - ${post_title}"
|
||||
body: "This is a test post created by the API testing framework"
|
||||
userId: "${user_id}"
|
||||
expected:
|
||||
status: 201
|
||||
json_path:
|
||||
- path: "$.id"
|
||||
operator: "greater_than"
|
||||
value: 100
|
||||
- path: "$.userId"
|
||||
value: "${user_id}"
|
||||
operator: "equals"
|
||||
extract:
|
||||
new_post_id:
|
||||
type: "json_path"
|
||||
path: "$.id"
|
||||
|
||||
- name: "Update created post"
|
||||
method: PUT
|
||||
url: "/posts/${new_post_id}"
|
||||
headers:
|
||||
Content-Type: "application/json"
|
||||
json:
|
||||
id: "${new_post_id}"
|
||||
title: "Updated Test Post"
|
||||
body: "This post has been updated"
|
||||
userId: "${user_id}"
|
||||
expected:
|
||||
status: 200
|
||||
json_path:
|
||||
- path: "$.title"
|
||||
value: "Updated Test Post"
|
||||
operator: "equals"
|
||||
|
||||
- name: "Delete created post"
|
||||
method: DELETE
|
||||
url: "/posts/${new_post_id}"
|
||||
expected:
|
||||
status: 200
|
||||
|
||||
- name: "Verify post was deleted"
|
||||
method: GET
|
||||
url: "/posts/${new_post_id}"
|
||||
expected:
|
||||
status: 404
|
||||
|
||||
- name: "Test with query parameters"
|
||||
method: GET
|
||||
url: "/posts"
|
||||
params:
|
||||
userId: "${user_id}"
|
||||
expected:
|
||||
status: 200
|
||||
json_path:
|
||||
- path: "$[*].userId"
|
||||
value: "${user_id}"
|
||||
operator: "equals"
|
||||
|
||||
- name: "Test response time validation"
|
||||
method: GET
|
||||
url: "/posts/1"
|
||||
expected:
|
||||
status: 200
|
||||
max_response_time: 2.0
|
||||
json_path:
|
||||
- path: "$.id"
|
||||
value: 1
|
||||
operator: "equals"
|
||||
|
||||
- name: "Test header validation"
|
||||
method: GET
|
||||
url: "/posts/1"
|
||||
expected:
|
||||
status: 200
|
||||
headers:
|
||||
Content-Type: "application/json; charset=utf-8"
|
||||
|
||||
- name: "Test multiple body contains assertions"
|
||||
method: GET
|
||||
url: "/posts/1"
|
||||
expected:
|
||||
status: 200
|
||||
body_contains:
|
||||
- "userId"
|
||||
- "title"
|
||||
- "body"
|
||||
|
||||
- name: "Test regex extraction from response"
|
||||
method: GET
|
||||
url: "/posts/1"
|
||||
extract:
|
||||
content_type:
|
||||
type: "header"
|
||||
path: "Content-Type"
|
||||
expected:
|
||||
status: 200
|
||||
@@ -0,0 +1,121 @@
|
||||
import re
|
||||
from typing import Any, Dict, Optional
|
||||
from jsonpath_ng import parse as jsonpath_parse
|
||||
|
||||
|
||||
class VariableContext:
|
||||
"""Manages variable extraction and substitution for request chaining."""
|
||||
|
||||
def __init__(self):
|
||||
self.variables: Dict[str, Any] = {}
|
||||
|
||||
def extract_variables(self, response, extractions: Dict[str, Dict[str, str]]) -> None:
|
||||
"""Extract variables from HTTP response based on extraction rules.
|
||||
|
||||
Args:
|
||||
response: HTTP response object
|
||||
extractions: Dict mapping variable names to extraction configs
|
||||
Format: {var_name: {"type": "json_path|header|regex", "path": "..."}}
|
||||
"""
|
||||
if not extractions:
|
||||
return
|
||||
|
||||
for var_name, config in extractions.items():
|
||||
extraction_type = config.get("type", "json_path")
|
||||
path = config.get("path", "")
|
||||
|
||||
try:
|
||||
if extraction_type == "json_path":
|
||||
self._extract_from_json(var_name, response, path)
|
||||
elif extraction_type == "header":
|
||||
self._extract_from_header(var_name, response, path)
|
||||
elif extraction_type == "regex":
|
||||
self._extract_from_regex(var_name, response, path)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to extract variable '{var_name}': {e}")
|
||||
|
||||
def _extract_from_json(self, var_name: str, response, path: str) -> None:
|
||||
"""Extract variable from JSON response using JSONPath."""
|
||||
try:
|
||||
json_data = response.json()
|
||||
jsonpath_expr = jsonpath_parse(path)
|
||||
matches = jsonpath_expr.find(json_data)
|
||||
if matches:
|
||||
self.variables[var_name] = matches[0].value
|
||||
except Exception as e:
|
||||
raise ValueError(f"JSON path extraction failed: {e}")
|
||||
|
||||
def _extract_from_header(self, var_name: str, response, path: str) -> None:
|
||||
"""Extract variable from response header."""
|
||||
header_value = response.headers.get(path)
|
||||
if header_value is not None:
|
||||
self.variables[var_name] = header_value
|
||||
else:
|
||||
raise ValueError(f"Header '{path}' not found")
|
||||
|
||||
def _extract_from_regex(self, var_name: str, response, path: str) -> None:
|
||||
"""Extract variable from response text using regex."""
|
||||
match = re.search(path, response.text)
|
||||
if match:
|
||||
self.variables[var_name] = match.group(1) if match.groups() else match.group(0)
|
||||
else:
|
||||
raise ValueError(f"Regex pattern '{path}' not found")
|
||||
|
||||
def substitute_variables(self, text: str) -> str:
|
||||
"""Replace variable placeholders in text with actual values.
|
||||
|
||||
Args:
|
||||
text: String containing variable placeholders like ${variable_name}
|
||||
|
||||
Returns:
|
||||
String with variables substituted
|
||||
"""
|
||||
if not isinstance(text, str):
|
||||
return text
|
||||
|
||||
def replace_var(match):
|
||||
var_name = match.group(1)
|
||||
if var_name in self.variables:
|
||||
return str(self.variables[var_name])
|
||||
return match.group(0) # Return original if variable not found
|
||||
|
||||
return re.sub(r'\$\{([^}]+)\}', replace_var, text)
|
||||
|
||||
def substitute_in_dict(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Recursively substitute variables in dictionary values."""
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
|
||||
result = {}
|
||||
for key, value in data.items():
|
||||
if isinstance(value, str):
|
||||
result[key] = self.substitute_variables(value)
|
||||
elif isinstance(value, dict):
|
||||
result[key] = self.substitute_in_dict(value)
|
||||
elif isinstance(value, list):
|
||||
result[key] = self.substitute_in_list(value)
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
def substitute_in_list(self, data: list) -> list:
|
||||
"""Recursively substitute variables in list items."""
|
||||
result = []
|
||||
for item in data:
|
||||
if isinstance(item, str):
|
||||
result.append(self.substitute_variables(item))
|
||||
elif isinstance(item, dict):
|
||||
result.append(self.substitute_in_dict(item))
|
||||
elif isinstance(item, list):
|
||||
result.append(self.substitute_in_list(item))
|
||||
else:
|
||||
result.append(item)
|
||||
return result
|
||||
|
||||
def get_variable(self, name: str) -> Optional[Any]:
|
||||
"""Get variable value by name."""
|
||||
return self.variables.get(name)
|
||||
|
||||
def set_variable(self, name: str, value: Any) -> None:
|
||||
"""Set variable value."""
|
||||
self.variables[name] = value
|
||||
@@ -0,0 +1,670 @@
|
||||
I'll create a production-ready FastAPI bookstore API with SQLite database and comprehensive CRUD operations. Here's the complete implementation:
|
||||
|
||||
## requirements.txt
|
||||
```txt
|
||||
fastapi>=0.104.0
|
||||
uvicorn[standard]>=0.24.0
|
||||
sqlalchemy>=2.0.0
|
||||
pydantic>=2.0.0
|
||||
python-multipart
|
||||
```
|
||||
|
||||
## database.py
|
||||
```python
|
||||
"""Database configuration and session management."""
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
import os
|
||||
|
||||
# Database URL
|
||||
SQLALCHEMY_DATABASE_URL = "sqlite:///./bookstore.db"
|
||||
|
||||
# Create engine with connection pooling and foreign key support
|
||||
engine = create_engine(
|
||||
SQLALCHEMY_DATABASE_URL,
|
||||
connect_args={
|
||||
"check_same_thread": False, # Allow multiple threads
|
||||
"timeout": 20, # 20 second timeout
|
||||
},
|
||||
pool_pre_ping=True, # Verify connections before use
|
||||
echo=False # Set to True for SQL debugging
|
||||
)
|
||||
|
||||
# Session factory
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
|
||||
# Base class for models
|
||||
Base = declarative_base()
|
||||
|
||||
def get_db():
|
||||
"""Dependency for database sessions."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
```
|
||||
|
||||
## models.py
|
||||
```python
|
||||
"""SQLAlchemy database models."""
|
||||
|
||||
from sqlalchemy import Column, Integer, String, Float, DateTime, Text
|
||||
from sqlalchemy.sql import func
|
||||
from database import Base
|
||||
import re
|
||||
|
||||
class Book(Base):
|
||||
"""Book model with validation and indexing."""
|
||||
|
||||
__tablename__ = "books"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
title = Column(String(255), nullable=False, index=True)
|
||||
author = Column(String(255), nullable=False, index=True)
|
||||
isbn = Column(String(17), unique=True, nullable=False, index=True) # ISBN-13 format
|
||||
price = Column(Float, nullable=False)
|
||||
publication_year = Column(Integer, nullable=True)
|
||||
genre = Column(String(100), nullable=True, index=True)
|
||||
description = Column(Text, nullable=True)
|
||||
stock_quantity = Column(Integer, default=0, nullable=False)
|
||||
created_at = Column(DateTime(timezone=True), server_default=func.now())
|
||||
updated_at = Column(DateTime(timezone=True), onupdate=func.now())
|
||||
|
||||
def __repr__(self):
|
||||
return f"<Book(id={self.id}, title='{self.title}', author='{self.author}')>"
|
||||
|
||||
@staticmethod
|
||||
def validate_isbn(isbn: str) -> bool:
|
||||
"""Validate ISBN-10 or ISBN-13 format."""
|
||||
if not isbn:
|
||||
return False
|
||||
|
||||
# Remove hyphens and spaces
|
||||
isbn_clean = re.sub(r'[-\s]', '', isbn)
|
||||
|
||||
# Check ISBN-10 or ISBN-13
|
||||
if len(isbn_clean) == 10:
|
||||
return re.match(r'^\d{9}[\dX]$', isbn_clean) is not None
|
||||
elif len(isbn_clean) == 13:
|
||||
return re.match(r'^\d{13}$', isbn_clean) is not None
|
||||
|
||||
return False
|
||||
```
|
||||
|
||||
## schemas.py
|
||||
```python
|
||||
"""Pydantic schemas for request/response validation."""
|
||||
|
||||
from pydantic import BaseModel, Field, validator
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
from models import Book as BookModel
|
||||
|
||||
class BookBase(BaseModel):
|
||||
"""Base schema for book data."""
|
||||
title: str = Field(..., min_length=1, max_length=255, description="Book title")
|
||||
author: str = Field(..., min_length=1, max_length=255, description="Book author")
|
||||
isbn: str = Field(..., min_length=10, max_length=17, description="ISBN-10 or ISBN-13")
|
||||
price: float = Field(..., gt=0, description="Book price (must be positive)")
|
||||
publication_year: Optional[int] = Field(None, ge=1000, le=2100, description="Publication year")
|
||||
genre: Optional[str] = Field(None, max_length=100, description="Book genre")
|
||||
description: Optional[str] = Field(None, description="Book description")
|
||||
stock_quantity: int = Field(0, ge=0, description="Stock quantity (non-negative)")
|
||||
|
||||
@validator('isbn')
|
||||
def validate_isbn(cls, v):
|
||||
"""Validate ISBN format."""
|
||||
if not BookModel.validate_isbn(v):
|
||||
raise ValueError('Invalid ISBN format. Must be valid ISBN-10 or ISBN-13')
|
||||
return v
|
||||
|
||||
@validator('title', 'author')
|
||||
def validate_strings(cls, v):
|
||||
"""Validate string fields are not empty after stripping."""
|
||||
if not v or not v.strip():
|
||||
raise ValueError('Field cannot be empty or whitespace only')
|
||||
return v.strip()
|
||||
|
||||
class BookCreate(BookBase):
|
||||
"""Schema for creating a new book."""
|
||||
pass
|
||||
|
||||
class BookUpdate(BaseModel):
|
||||
"""Schema for updating a book (all fields optional)."""
|
||||
title: Optional[str] = Field(None, min_length=1, max_length=255)
|
||||
author: Optional[str] = Field(None, min_length=1, max_length=255)
|
||||
isbn: Optional[str] = Field(None, min_length=10, max_length=17)
|
||||
price: Optional[float] = Field(None, gt=0)
|
||||
publication_year: Optional[int] = Field(None, ge=1000, le=2100)
|
||||
genre: Optional[str] = Field(None, max_length=100)
|
||||
description: Optional[str] = None
|
||||
stock_quantity: Optional[int] = Field(None, ge=0)
|
||||
|
||||
@validator('isbn', allow_reuse=True)
|
||||
def validate_isbn(cls, v):
|
||||
if v is not None and not BookModel.validate_isbn(v):
|
||||
raise ValueError('Invalid ISBN format')
|
||||
return v
|
||||
|
||||
@validator('title', 'author', allow_reuse=True)
|
||||
def validate_strings(cls, v):
|
||||
if v is not None and (not v or not v.strip()):
|
||||
raise ValueError('Field cannot be empty or whitespace only')
|
||||
return v.strip() if v else v
|
||||
|
||||
class BookResponse(BookBase):
|
||||
"""Schema for book responses."""
|
||||
id: int
|
||||
created_at: datetime
|
||||
updated_at: Optional[datetime] = None
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
class BookListResponse(BaseModel):
|
||||
"""Schema for paginated book list responses."""
|
||||
books: list[BookResponse]
|
||||
total: int
|
||||
page: int
|
||||
per_page: int
|
||||
total_pages: int
|
||||
|
||||
class ErrorResponse(BaseModel):
|
||||
"""Schema for error responses."""
|
||||
error: str
|
||||
detail: Optional[str] = None
|
||||
status_code: int
|
||||
```
|
||||
|
||||
## crud.py
|
||||
```python
|
||||
"""CRUD operations for books."""
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import or_, and_
|
||||
from models import Book
|
||||
from schemas import BookCreate, BookUpdate
|
||||
from typing import Optional, List
|
||||
|
||||
class BookCRUD:
|
||||
"""Book CRUD operations with error handling."""
|
||||
|
||||
@staticmethod
|
||||
def create_book(db: Session, book: BookCreate) -> Book:
|
||||
"""Create a new book."""
|
||||
try:
|
||||
db_book = Book(**book.dict())
|
||||
db.add(db_book)
|
||||
db.commit()
|
||||
db.refresh(db_book)
|
||||
return db_book
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def get_book(db: Session, book_id: int) -> Optional[Book]:
|
||||
"""Get a book by ID."""
|
||||
return db.query(Book).filter(Book.id == book_id).first()
|
||||
|
||||
@staticmethod
|
||||
def get_book_by_isbn(db: Session, isbn: str) -> Optional[Book]:
|
||||
"""Get a book by ISBN."""
|
||||
return db.query(Book).filter(Book.isbn == isbn).first()
|
||||
|
||||
@staticmethod
|
||||
def get_books(
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
title: Optional[str] = None,
|
||||
author: Optional[str] = None,
|
||||
genre: Optional[str] = None,
|
||||
min_price: Optional[float] = None,
|
||||
max_price: Optional[float] = None
|
||||
) -> tuple[List[Book], int]:
|
||||
"""Get books with optional filtering and pagination."""
|
||||
query = db.query(Book)
|
||||
|
||||
# Apply filters
|
||||
filters = []
|
||||
if title:
|
||||
filters.append(Book.title.ilike(f"%{title}%"))
|
||||
if author:
|
||||
filters.append(Book.author.ilike(f"%{author}%"))
|
||||
if genre:
|
||||
filters.append(Book.genre.ilike(f"%{genre}%"))
|
||||
if min_price is not None:
|
||||
filters.append(Book.price >= min_price)
|
||||
if max_price is not None:
|
||||
filters.append(Book.price <= max_price)
|
||||
|
||||
if filters:
|
||||
query = query.filter(and_(*filters))
|
||||
|
||||
# Get total count
|
||||
total = query.count()
|
||||
|
||||
# Apply pagination
|
||||
books = query.offset(skip).limit(limit).all()
|
||||
|
||||
return books, total
|
||||
|
||||
@staticmethod
|
||||
def update_book(db: Session, book_id: int, book_update: BookUpdate) -> Optional[Book]:
|
||||
"""Update a book by ID."""
|
||||
try:
|
||||
db_book = db.query(Book).filter(Book.id == book_id).first()
|
||||
if not db_book:
|
||||
return None
|
||||
|
||||
# Update only provided fields
|
||||
update_data = book_update.dict(exclude_unset=True)
|
||||
for field, value in update_data.items():
|
||||
setattr(db_book, field, value)
|
||||
|
||||
db.commit()
|
||||
db.refresh(db_book)
|
||||
return db_book
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def delete_book(db: Session, book_id: int) -> bool:
|
||||
"""Delete a book by ID."""
|
||||
try:
|
||||
db_book = db.query(Book).filter(Book.id == book_id).first()
|
||||
if not db_book:
|
||||
return False
|
||||
|
||||
db.delete(db_book)
|
||||
db.commit()
|
||||
return True
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def search_books(db: Session, query: str, limit: int = 50) -> List[Book]:
|
||||
"""Search books by title, author, or genre."""
|
||||
search_filter = or_(
|
||||
Book.title.ilike(f"%{query}%"),
|
||||
Book.author.ilike(f"%{query}%"),
|
||||
Book.genre.ilike(f"%{query}%")
|
||||
)
|
||||
return db.query(Book).filter(search_filter).limit(limit).all()
|
||||
```
|
||||
|
||||
## main.py
|
||||
```python
|
||||
"""FastAPI bookstore application."""
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Depends, Query, status
|
||||
from fastapi.responses import JSONResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
import math
|
||||
|
||||
from database import SessionLocal, engine, get_db
|
||||
from models import Base, Book
|
||||
from schemas import (
|
||||
BookCreate, BookUpdate, BookResponse, BookListResponse,
|
||||
ErrorResponse
|
||||
)
|
||||
from crud import BookCRUD
|
||||
|
||||
# Create database tables
|
||||
Base.metadata.create_all(bind=engine)
|
||||
|
||||
# Initialize FastAPI app
|
||||
app = FastAPI(
|
||||
title="Bookstore API",
|
||||
description="A comprehensive RESTful API for managing a bookstore inventory",
|
||||
version="1.0.0",
|
||||
docs_url="/docs",
|
||||
redoc_url="/redoc"
|
||||
)
|
||||
|
||||
# Initialize CRUD operations
|
||||
book_crud = BookCRUD()
|
||||
|
||||
@app.exception_handler(ValueError)
|
||||
async def value_error_handler(request, exc):
|
||||
"""Handle validation errors."""
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
content={"error": "Validation Error", "detail": str(exc)}
|
||||
)
|
||||
|
||||
@app.exception_handler(IntegrityError)
|
||||
async def integrity_error_handler(request, exc):
|
||||
"""Handle database integrity errors."""
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
content={"error": "Conflict", "detail": "Book with this ISBN already exists"}
|
||||
)
|
||||
|
||||
@app.get("/", tags=["Root"])
|
||||
async def root():
|
||||
"""Root endpoint with API information."""
|
||||
return {
|
||||
"message": "Bookstore API",
|
||||
"version": "1.0.0",
|
||||
"docs": "/docs",
|
||||
"endpoints": {
|
||||
"books": "/books",
|
||||
"search": "/books/search"
|
||||
}
|
||||
}
|
||||
|
||||
@app.post(
|
||||
"/books",
|
||||
response_model=BookResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
tags=["Books"],
|
||||
summary="Create a new book"
|
||||
)
|
||||
async def create_book(book: BookCreate, db: Session = Depends(get_db)):
|
||||
"""
|
||||
Create a new book in the inventory.
|
||||
|
||||
- **title**: Book title (required)
|
||||
- **author**: Book author (required)
|
||||
- **isbn**: Valid ISBN-10 or ISBN-13 (required, unique)
|
||||
- **price**: Book price in USD (required, positive)
|
||||
- **publication_year**: Year of publication (optional)
|
||||
- **genre**: Book genre (optional)
|
||||
- **description**: Book description (optional)
|
||||
- **stock_quantity**: Number in stock (default: 0)
|
||||
"""
|
||||
try:
|
||||
# Check if ISBN already exists
|
||||
if book_crud.get_book_by_isbn(db, book.isbn):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Book with this ISBN already exists"
|
||||
)
|
||||
|
||||
db_book = book_crud.create_book(db, book)
|
||||
return db_book
|
||||
except IntegrityError:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Book with this ISBN already exists"
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Failed to create book"
|
||||
)
|
||||
|
||||
@app.get(
|
||||
"/books",
|
||||
response_model=BookListResponse,
|
||||
tags=["Books"],
|
||||
summary="Get all books with optional filtering and pagination"
|
||||
)
|
||||
async def get_books(
|
||||
page: int = Query(1, ge=1, description="Page number"),
|
||||
per_page: int = Query(10, ge=1, le=100, description="Books per page"),
|
||||
title: str = Query(None, description="Filter by title (partial match)"),
|
||||
author: str = Query(None, description="Filter by author (partial match)"),
|
||||
genre: str = Query(None, description="Filter by genre (partial match)"),
|
||||
min_price: float = Query(None, ge=0, description="Minimum price filter"),
|
||||
max_price: float = Query(None, ge=0, description="Maximum price filter"),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Retrieve books with optional filtering and pagination.
|
||||
|
||||
Supports filtering by:
|
||||
- Title (partial text search)
|
||||
- Author (partial text search)
|
||||
- Genre (partial text search)
|
||||
- Price range (min_price, max_price)
|
||||
"""
|
||||
try:
|
||||
skip = (page - 1) * per_page
|
||||
|
||||
books, total = book_crud.get_books(
|
||||
db=db,
|
||||
skip=skip,
|
||||
limit=per_page,
|
||||
title=title,
|
||||
author=author,
|
||||
genre=genre,
|
||||
min_price=min_price,
|
||||
max_price=max_price
|
||||
)
|
||||
|
||||
total_pages = math.ceil(total / per_page) if total > 0 else 1
|
||||
|
||||
return BookListResponse(
|
||||
books=books,
|
||||
total=total,
|
||||
page=page,
|
||||
per_page=per_page,
|
||||
total_pages=total_pages
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Failed to retrieve books"
|
||||
)
|
||||
|
||||
@app.get(
|
||||
"/books/{book_id}",
|
||||
response_model=BookResponse,
|
||||
tags=["Books"],
|
||||
summary="Get a specific book by ID"
|
||||
)
|
||||
async def get_book(book_id: int, db: Session = Depends(get_db)):
|
||||
"""Get a specific book by its ID."""
|
||||
try:
|
||||
book = book_crud.get_book(db, book_id)
|
||||
if not book:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Book not found"
|
||||
)
|
||||
return book
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Failed to retrieve book"
|
||||
)
|
||||
|
||||
@app.put(
|
||||
"/books/{book_id}",
|
||||
response_model=BookResponse,
|
||||
tags=["Books"],
|
||||
summary="Update a book completely"
|
||||
)
|
||||
async def update_book_put(
|
||||
book_id: int,
|
||||
book: BookCreate,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Update all fields of a book (complete replacement)."""
|
||||
try:
|
||||
# Check if book exists
|
||||
existing_book = book_crud.get_book(db, book_id)
|
||||
if not existing_book:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Book not found"
|
||||
)
|
||||
|
||||
# Check ISBN uniqueness (if different from current)
|
||||
if book.isbn != existing_book.isbn:
|
||||
isbn_book = book_crud.get_book_by_isbn(db, book.isbn)
|
||||
if isbn_book:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Book with this ISBN already exists"
|
||||
)
|
||||
|
||||
# Convert BookCreate to BookUpdate for consistency
|
||||
book_update = BookUpdate(**book.dict())
|
||||
updated_book = book_crud.update_book(db, book_id, book_update)
|
||||
return updated_book
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Failed to update book"
|
||||
)
|
||||
|
||||
@app.patch(
|
||||
"/books/{book_id}",
|
||||
response_model=BookResponse,
|
||||
tags=["Books"],
|
||||
summary="Partially update a book"
|
||||
)
|
||||
async def update_book_patch(
|
||||
book_id: int,
|
||||
book: BookUpdate,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Update specific fields of a book (partial update)."""
|
||||
try:
|
||||
# Check if book exists
|
||||
existing_book = book_crud.get_book(db, book_id)
|
||||
if not existing_book:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Book not found"
|
||||
)
|
||||
|
||||
# Check ISBN uniqueness (if being updated)
|
||||
if book.isbn and book.isbn != existing_book.isbn:
|
||||
isbn_book = book_crud.get_book_by_isbn(db, book.isbn)
|
||||
if isbn_book:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Book with this ISBN already exists"
|
||||
)
|
||||
|
||||
updated_book = book_crud.update_book(db, book_id, book)
|
||||
return updated_book
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Failed to update book"
|
||||
)
|
||||
|
||||
@app.delete(
|
||||
"/books/{book_id}",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
tags=["Books"],
|
||||
summary="Delete a book"
|
||||
)
|
||||
async def delete_book(book_id: int, db: Session = Depends(get_db)):
|
||||
"""Delete a book from the inventory."""
|
||||
try:
|
||||
deleted = book_crud.delete_book(db, book_id)
|
||||
if not deleted:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Book not found"
|
||||
)
|
||||
return None
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Failed to delete book"
|
||||
)
|
||||
|
||||
@app.get(
|
||||
"/books/search/{query}",
|
||||
response_model=list[BookResponse],
|
||||
tags=["Books"],
|
||||
summary="Search books by title, author, or genre"
|
||||
)
|
||||
async def search_books(
|
||||
query: str,
|
||||
limit: int = Query(50, ge=1, le=100, description="Maximum results to return"),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Search books by query string across title, author, and genre fields.
|
||||
|
||||
Returns books where the query matches any part of:
|
||||
- Book title
|
||||
- Author name
|
||||
- Genre
|
||||
"""
|
||||
try:
|
||||
if len(query.strip()) < 2:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail="Search query must be at least 2 characters long"
|
||||
)
|
||||
|
||||
books = book_crud.search_books(db, query.strip(), limit)
|
||||
return books
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Search failed"
|
||||
)
|
||||
|
||||
@app.get("/health", tags=["Health"])
|
||||
async def health_check():
|
||||
"""Health check endpoint."""
|
||||
return {"status": "healthy", "service": "bookstore-api"}
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(
|
||||
"main:app",
|
||||
host="0.0.0.0",
|
||||
port=8000,
|
||||
reload=True,
|
||||
log_level="info"
|
||||
)
|
||||
```
|
||||
|
||||
## Usage Instructions
|
||||
|
||||
1. **Install dependencies**:
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
2. **Run the application**:
|
||||
```bash
|
||||
uvicorn main:app --reload
|
||||
```
|
||||
|
||||
3. **Access the API**:
|
||||
- API: http://localhost:8000
|
||||
- Interactive docs: http://localhost:8000/docs
|
||||
- ReDoc: http://localhost:8000/redoc
|
||||
|
||||
## Key Features
|
||||
|
||||
- **Complete CRUD operations** with proper HTTP methods
|
||||
- **Advanced filtering and pagination** for book listings
|
||||
- **ISBN validation** for both ISBN-10 and ISBN-13 formats
|
||||
- **Comprehensive error handling** with appropriate HTTP status codes
|
||||
- **Database connection pooling** and session management
|
||||
- **Input validation** using Pydantic schemas
|
||||
- **Search functionality** across multiple fields
|
||||
- **Auto-generated OpenAPI documentation**
|
||||
- **Production-ready error handling** and logging setup
|
||||
|
||||
The API provides robust data validation, proper error responses, and follows REST principles with comprehensive documentation.
|
||||
@@ -0,0 +1,120 @@
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Dict, List, Optional
|
||||
import uvicorn
|
||||
|
||||
app = FastAPI(
|
||||
title="Bookstore API",
|
||||
description="A simple REST API for managing books in a bookstore",
|
||||
version="1.0.0"
|
||||
)
|
||||
|
||||
# Pydantic models
|
||||
class Book(BaseModel):
|
||||
id: int = Field(..., description="Unique identifier for the book")
|
||||
title: str = Field(..., min_length=1, description="Title of the book")
|
||||
author: str = Field(..., min_length=1, description="Author of the book")
|
||||
price: float = Field(..., gt=0, description="Price of the book in dollars")
|
||||
isbn: str = Field(..., min_length=10, max_length=17, description="ISBN of the book")
|
||||
|
||||
class BookCreate(BaseModel):
|
||||
title: str = Field(..., min_length=1, description="Title of the book")
|
||||
author: str = Field(..., min_length=1, description="Author of the book")
|
||||
price: float = Field(..., gt=0, description="Price of the book in dollars")
|
||||
isbn: str = Field(..., min_length=10, max_length=17, description="ISBN of the book")
|
||||
|
||||
class BookUpdate(BaseModel):
|
||||
title: Optional[str] = Field(None, min_length=1, description="Title of the book")
|
||||
author: Optional[str] = Field(None, min_length=1, description="Author of the book")
|
||||
price: Optional[float] = Field(None, gt=0, description="Price of the book in dollars")
|
||||
isbn: Optional[str] = Field(None, min_length=10, max_length=17, description="ISBN of the book")
|
||||
|
||||
# In-memory storage
|
||||
books_db: Dict[int, Book] = {}
|
||||
next_book_id = 1
|
||||
|
||||
# Sample data
|
||||
sample_books = [
|
||||
{"title": "The Great Gatsby", "author": "F. Scott Fitzgerald", "price": 12.99, "isbn": "978-0-7432-7356-5"},
|
||||
{"title": "To Kill a Mockingbird", "author": "Harper Lee", "price": 14.99, "isbn": "978-0-06-112008-4"},
|
||||
{"title": "1984", "author": "George Orwell", "price": 13.99, "isbn": "978-0-452-28423-4"}
|
||||
]
|
||||
|
||||
# Initialize sample data
|
||||
for book_data in sample_books:
|
||||
book = Book(id=next_book_id, **book_data)
|
||||
books_db[next_book_id] = book
|
||||
next_book_id += 1
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"message": "Welcome to the Bookstore API", "docs": "/docs"}
|
||||
|
||||
@app.get("/books", response_model=List[Book])
|
||||
async def get_all_books():
|
||||
"""Retrieve all books from the bookstore."""
|
||||
return list(books_db.values())
|
||||
|
||||
@app.get("/books/{book_id}", response_model=Book)
|
||||
async def get_book(book_id: int):
|
||||
"""Retrieve a specific book by its ID."""
|
||||
if book_id not in books_db:
|
||||
raise HTTPException(status_code=404, detail="Book not found")
|
||||
return books_db[book_id]
|
||||
|
||||
@app.post("/books", response_model=Book, status_code=201)
|
||||
async def create_book(book: BookCreate):
|
||||
"""Create a new book in the bookstore."""
|
||||
global next_book_id
|
||||
|
||||
# Check if ISBN already exists
|
||||
for existing_book in books_db.values():
|
||||
if existing_book.isbn == book.isbn:
|
||||
raise HTTPException(status_code=400, detail="Book with this ISBN already exists")
|
||||
|
||||
new_book = Book(id=next_book_id, **book.dict())
|
||||
books_db[next_book_id] = new_book
|
||||
next_book_id += 1
|
||||
|
||||
return new_book
|
||||
|
||||
@app.put("/books/{book_id}", response_model=Book)
|
||||
async def update_book(book_id: int, book_update: BookUpdate):
|
||||
"""Update an existing book by its ID."""
|
||||
if book_id not in books_db:
|
||||
raise HTTPException(status_code=404, detail="Book not found")
|
||||
|
||||
existing_book = books_db[book_id]
|
||||
|
||||
# Check if ISBN already exists for another book
|
||||
if book_update.isbn:
|
||||
for bid, existing in books_db.items():
|
||||
if bid != book_id and existing.isbn == book_update.isbn:
|
||||
raise HTTPException(status_code=400, detail="Book with this ISBN already exists")
|
||||
|
||||
# Update only provided fields
|
||||
update_data = book_update.dict(exclude_unset=True)
|
||||
updated_book_data = existing_book.dict()
|
||||
updated_book_data.update(update_data)
|
||||
|
||||
updated_book = Book(**updated_book_data)
|
||||
books_db[book_id] = updated_book
|
||||
|
||||
return updated_book
|
||||
|
||||
@app.delete("/books/{book_id}", status_code=204)
|
||||
async def delete_book(book_id: int):
|
||||
"""Delete a book from the bookstore by its ID."""
|
||||
if book_id not in books_db:
|
||||
raise HTTPException(status_code=404, detail="Book not found")
|
||||
|
||||
del books_db[book_id]
|
||||
return None
|
||||
|
||||
# Health check endpoint
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
return {"status": "healthy", "total_books": len(books_db)}
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
@@ -0,0 +1,3 @@
|
||||
fastapi>=0.104.0
|
||||
uvicorn>=0.24.0
|
||||
pydantic>=2.0.0
|
||||
@@ -0,0 +1,451 @@
|
||||
I'll create a production-ready WebSocket chat server with proper error handling and connection management. Here are the files:
|
||||
|
||||
## chat_server.py
|
||||
```python
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
WebSocket Chat Server
|
||||
A production-ready chat server supporting multiple concurrent clients.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import signal
|
||||
import sys
|
||||
import websockets
|
||||
from datetime import datetime
|
||||
from typing import Set, Dict, Any
|
||||
from websockets.server import WebSocketServerProtocol
|
||||
from websockets.exceptions import ConnectionClosed, WebSocketException
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ChatServer:
|
||||
"""WebSocket chat server managing multiple client connections."""
|
||||
|
||||
def __init__(self, host: str = "localhost", port: int = 8765):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.clients: Set[WebSocketServerProtocol] = set()
|
||||
self.server = None
|
||||
self._shutdown_event = asyncio.Event()
|
||||
|
||||
async def register_client(self, websocket: WebSocketServerProtocol) -> None:
|
||||
"""Register a new client connection."""
|
||||
self.clients.add(websocket)
|
||||
client_info = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
|
||||
logger.info(f"Client connected: {client_info} (Total: {len(self.clients)})")
|
||||
|
||||
# Send welcome message to new client
|
||||
welcome_msg = self._create_message(
|
||||
"system",
|
||||
f"Welcome to the chat! {len(self.clients)} user(s) online."
|
||||
)
|
||||
await self._safe_send(websocket, welcome_msg)
|
||||
|
||||
async def unregister_client(self, websocket: WebSocketServerProtocol) -> None:
|
||||
"""Unregister a client connection."""
|
||||
self.clients.discard(websocket)
|
||||
client_info = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
|
||||
logger.info(f"Client disconnected: {client_info} (Total: {len(self.clients)})")
|
||||
|
||||
async def broadcast_message(self, message: str, sender: WebSocketServerProtocol) -> None:
|
||||
"""Broadcast a message to all connected clients except the sender."""
|
||||
if not self.clients:
|
||||
return
|
||||
|
||||
# Create structured message
|
||||
chat_message = self._create_message("user", message, sender.remote_address[0])
|
||||
|
||||
# Send to all clients except sender
|
||||
recipients = self.clients - {sender}
|
||||
if recipients:
|
||||
await asyncio.gather(
|
||||
*[self._safe_send(client, chat_message) for client in recipients],
|
||||
return_exceptions=True
|
||||
)
|
||||
logger.info(f"Broadcasted message to {len(recipients)} clients")
|
||||
|
||||
def _create_message(self, msg_type: str, content: str, sender: str = "system") -> str:
|
||||
"""Create a structured JSON message."""
|
||||
return json.dumps({
|
||||
"type": msg_type,
|
||||
"sender": sender,
|
||||
"content": content,
|
||||
"timestamp": datetime.now().isoformat()
|
||||
})
|
||||
|
||||
async def _safe_send(self, websocket: WebSocketServerProtocol, message: str) -> None:
|
||||
"""Safely send a message to a client with error handling."""
|
||||
try:
|
||||
await websocket.send(message)
|
||||
except ConnectionClosed:
|
||||
logger.debug(f"Connection closed for {websocket.remote_address}")
|
||||
await self.unregister_client(websocket)
|
||||
except WebSocketException as e:
|
||||
logger.warning(f"WebSocket error sending to {websocket.remote_address}: {e}")
|
||||
await self.unregister_client(websocket)
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error sending to {websocket.remote_address}: {e}")
|
||||
await self.unregister_client(websocket)
|
||||
|
||||
async def handle_client(self, websocket: WebSocketServerProtocol, path: str) -> None:
|
||||
"""Handle a client connection throughout its lifecycle."""
|
||||
await self.register_client(websocket)
|
||||
|
||||
try:
|
||||
async for message in websocket:
|
||||
try:
|
||||
# Validate message
|
||||
if not isinstance(message, str):
|
||||
continue
|
||||
|
||||
message = message.strip()
|
||||
if not message or len(message) > 1000: # Message length limit
|
||||
continue
|
||||
|
||||
# Broadcast to other clients
|
||||
await self.broadcast_message(message, websocket)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Invalid JSON from {websocket.remote_address}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing message from {websocket.remote_address}: {e}")
|
||||
break
|
||||
|
||||
except ConnectionClosed:
|
||||
logger.debug(f"Client {websocket.remote_address} disconnected normally")
|
||||
except WebSocketException as e:
|
||||
logger.warning(f"WebSocket error for {websocket.remote_address}: {e}")
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error handling client {websocket.remote_address}: {e}")
|
||||
finally:
|
||||
await self.unregister_client(websocket)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the WebSocket server."""
|
||||
logger.info(f"Starting chat server on {self.host}:{self.port}")
|
||||
|
||||
self.server = await websockets.serve(
|
||||
self.handle_client,
|
||||
self.host,
|
||||
self.port,
|
||||
ping_interval=30, # Send ping every 30 seconds
|
||||
ping_timeout=10, # Wait 10 seconds for pong
|
||||
max_size=10**6, # Max message size: 1MB
|
||||
max_queue=32 # Max queued messages per connection
|
||||
)
|
||||
|
||||
logger.info(f"Chat server started successfully on ws://{self.host}:{self.port}")
|
||||
|
||||
# Wait for shutdown signal
|
||||
await self._shutdown_event.wait()
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the WebSocket server gracefully."""
|
||||
logger.info("Shutting down chat server...")
|
||||
|
||||
if self.server:
|
||||
self.server.close()
|
||||
await self.server.wait_closed()
|
||||
|
||||
# Close all client connections
|
||||
if self.clients:
|
||||
await asyncio.gather(
|
||||
*[client.close() for client in self.clients.copy()],
|
||||
return_exceptions=True
|
||||
)
|
||||
|
||||
self._shutdown_event.set()
|
||||
logger.info("Chat server stopped")
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Signal shutdown from sync context."""
|
||||
self._shutdown_event.set()
|
||||
|
||||
async def main():
|
||||
"""Main entry point."""
|
||||
server = ChatServer(host="0.0.0.0", port=8765)
|
||||
|
||||
# Setup signal handlers for graceful shutdown
|
||||
def signal_handler():
|
||||
logger.info("Received shutdown signal")
|
||||
server.shutdown()
|
||||
|
||||
if sys.platform != "win32":
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGTERM, signal.SIGINT):
|
||||
loop.add_signal_handler(sig, signal_handler)
|
||||
|
||||
try:
|
||||
await server.start()
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Received keyboard interrupt")
|
||||
except Exception as e:
|
||||
logger.error(f"Server error: {e}")
|
||||
finally:
|
||||
await server.stop()
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutdown complete.")
|
||||
```
|
||||
|
||||
## requirements.txt
|
||||
```txt
|
||||
websockets>=11.0.2
|
||||
```
|
||||
|
||||
## README.md
|
||||
```markdown
|
||||
# WebSocket Chat Server
|
||||
|
||||
A production-ready WebSocket chat server built with Python that supports multiple concurrent clients and real-time message broadcasting.
|
||||
|
||||
## Features
|
||||
|
||||
- **Multi-client Support**: Handle multiple concurrent connections
|
||||
- **Real-time Messaging**: Instant message broadcasting to all connected clients
|
||||
- **Graceful Connection Management**: Proper handling of client connections and disconnections
|
||||
- **Error Handling**: Robust error handling for network issues and malformed messages
|
||||
- **Structured Messages**: JSON-formatted messages with timestamps
|
||||
- **Resource Limits**: Message size limits and connection management
|
||||
- **Logging**: Comprehensive logging for monitoring and debugging
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.7+
|
||||
- websockets library
|
||||
|
||||
## Installation
|
||||
|
||||
1. Clone or download the project files
|
||||
2. Install dependencies:
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Starting the Server
|
||||
|
||||
```bash
|
||||
python chat_server.py
|
||||
```
|
||||
|
||||
The server will start on `ws://0.0.0.0:8765` by default.
|
||||
|
||||
### Configuration
|
||||
|
||||
Modify the server settings in the `main()` function:
|
||||
|
||||
```python
|
||||
server = ChatServer(host="0.0.0.0", port=8765)
|
||||
```
|
||||
|
||||
### Connecting Clients
|
||||
|
||||
Connect to the WebSocket server using any WebSocket client:
|
||||
|
||||
```javascript
|
||||
const ws = new WebSocket('ws://localhost:8765');
|
||||
|
||||
ws.onmessage = function(event) {
|
||||
const message = JSON.parse(event.data);
|
||||
console.log(`[${message.timestamp}] ${message.sender}: ${message.content}`);
|
||||
};
|
||||
|
||||
ws.send("Hello, everyone!");
|
||||
```
|
||||
|
||||
## Message Format
|
||||
|
||||
The server uses structured JSON messages:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "user|system",
|
||||
"sender": "client_ip|system",
|
||||
"content": "message_content",
|
||||
"timestamp": "2023-12-07T10:30:00.000000"
|
||||
}
|
||||
```
|
||||
|
||||
## Security Considerations
|
||||
|
||||
- **No Authentication**: This is a basic implementation without user authentication
|
||||
- **Rate Limiting**: Consider implementing rate limiting for production use
|
||||
- **Input Validation**: Messages are limited to 1000 characters
|
||||
- **Network Security**: Use WSS (WebSocket Secure) in production
|
||||
|
||||
## Production Deployment
|
||||
|
||||
For production deployment, consider:
|
||||
|
||||
1. **Reverse Proxy**: Use nginx or similar for SSL termination
|
||||
2. **Process Management**: Use systemd, supervisor, or container orchestration
|
||||
3. **Monitoring**: Implement health checks and metrics collection
|
||||
4. **Scaling**: Use Redis or similar for multi-server message broadcasting
|
||||
|
||||
## Graceful Shutdown
|
||||
|
||||
The server supports graceful shutdown via:
|
||||
- `Ctrl+C` (SIGINT)
|
||||
- `SIGTERM` signal (Unix-like systems)
|
||||
|
||||
## Logging
|
||||
|
||||
Logs include:
|
||||
- Client connection/disconnection events
|
||||
- Message broadcasting statistics
|
||||
- Error conditions and warnings
|
||||
- Server startup/shutdown events
|
||||
|
||||
## License
|
||||
|
||||
MIT License - Feel free to use and modify as needed.
|
||||
```
|
||||
|
||||
## client_test.html
|
||||
```html
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>WebSocket Chat Test Client</title>
|
||||
<style>
|
||||
body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
|
||||
#messages { border: 1px solid #ccc; height: 300px; overflow-y: auto; padding: 10px; margin: 10px 0; background-color: #f9f9f9; }
|
||||
#messageInput { width: 70%; padding: 10px; }
|
||||
#sendButton { width: 25%; padding: 10px; }
|
||||
.message { margin: 5px 0; }
|
||||
.system { color: #666; font-style: italic; }
|
||||
.user { color: #000; }
|
||||
.status { padding: 10px; margin: 10px 0; border-radius: 5px; }
|
||||
.connected { background-color: #d4edda; color: #155724; }
|
||||
.disconnected { background-color: #f8d7da; color: #721c24; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>WebSocket Chat Test Client</h1>
|
||||
|
||||
<div id="status" class="status disconnected">Disconnected</div>
|
||||
|
||||
<div>
|
||||
<button id="connectButton" onclick="connect()">Connect</button>
|
||||
<button id="disconnectButton" onclick="disconnect()" disabled>Disconnect</button>
|
||||
</div>
|
||||
|
||||
<div id="messages"></div>
|
||||
|
||||
<div>
|
||||
<input type="text" id="messageInput" placeholder="Type your message here..." disabled>
|
||||
<button id="sendButton" onclick="sendMessage()" disabled>Send</button>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
let ws = null;
|
||||
const messagesDiv = document.getElementById('messages');
|
||||
const messageInput = document.getElementById('messageInput');
|
||||
const statusDiv = document.getElementById('status');
|
||||
const connectButton = document.getElementById('connectButton');
|
||||
const disconnectButton = document.getElementById('disconnectButton');
|
||||
const sendButton = document.getElementById('sendButton');
|
||||
|
||||
function connect() {
|
||||
try {
|
||||
ws = new WebSocket('ws://localhost:8765');
|
||||
|
||||
ws.onopen = function(event) {
|
||||
updateStatus('Connected', true);
|
||||
addMessage('Connected to chat server', 'system');
|
||||
};
|
||||
|
||||
ws.onmessage = function(event) {
|
||||
try {
|
||||
const message = JSON.parse(event.data);
|
||||
const timestamp = new Date(message.timestamp).toLocaleTimeString();
|
||||
addMessage(`[${timestamp}] ${message.sender}: ${message.content}`, message.type);
|
||||
} catch (e) {
|
||||
addMessage(event.data, 'user');
|
||||
}
|
||||
};
|
||||
|
||||
ws.onclose = function(event) {
|
||||
updateStatus('Disconnected', false);
|
||||
addMessage('Disconnected from chat server', 'system');
|
||||
};
|
||||
|
||||
ws.onerror = function(error) {
|
||||
addMessage('WebSocket error: ' + error, 'system');
|
||||
};
|
||||
|
||||
} catch (error) {
|
||||
addMessage('Connection error: ' + error, 'system');
|
||||
}
|
||||
}
|
||||
|
||||
function disconnect() {
|
||||
if (ws) {
|
||||
ws.close();
|
||||
}
|
||||
}
|
||||
|
||||
function sendMessage() {
|
||||
const message = messageInput.value.trim();
|
||||
if (message && ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(message);
|
||||
addMessage(`You: ${message}`, 'user');
|
||||
messageInput.value = '';
|
||||
}
|
||||
}
|
||||
|
||||
function addMessage(text, type) {
|
||||
const messageElement = document.createElement('div');
|
||||
messageElement.className = `message ${type}`;
|
||||
messageElement.textContent = text;
|
||||
messagesDiv.appendChild(messageElement);
|
||||
messagesDiv.scrollTop = messagesDiv.scrollHeight;
|
||||
}
|
||||
|
||||
function updateStatus(text, connected) {
|
||||
statusDiv.textContent = text;
|
||||
statusDiv.className = `status ${connected ? 'connected' : 'disconnected'}`;
|
||||
|
||||
connectButton.disabled = connected;
|
||||
disconnectButton.disabled = !connected;
|
||||
messageInput.disabled = !connected;
|
||||
sendButton.disabled = !connected;
|
||||
}
|
||||
|
||||
// Send message on Enter key
|
||||
messageInput.addEventListener('keypress', function(e) {
|
||||
if (e.key === 'Enter') {
|
||||
sendMessage();
|
||||
}
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
```
|
||||
|
||||
This implementation provides a robust, production-ready WebSocket chat server with:
|
||||
|
||||
1. **Proper connection management** with automatic cleanup
|
||||
2. **Error handling** for network issues and malformed data
|
||||
3. **Structured JSON messaging** with timestamps
|
||||
4. **Resource limits** to prevent abuse
|
||||
5. **Graceful shutdown** support
|
||||
6. **Comprehensive logging** for monitoring
|
||||
7. **Test client** for easy testing
|
||||
|
||||
The server can handle multiple concurrent clients and broadcasts messages efficiently while maintaining connection stability.
|
||||
@@ -0,0 +1,423 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>WebSocket Chat</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: Arial, sans-serif;
|
||||
margin: 0;
|
||||
padding: 20px;
|
||||
background-color: #f5f5f5;
|
||||
}
|
||||
|
||||
.chat-container {
|
||||
max-width: 800px;
|
||||
margin: 0 auto;
|
||||
background: white;
|
||||
border-radius: 10px;
|
||||
box-shadow: 0 2px 10px rgba(0,0,0,0.1);
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.chat-header {
|
||||
background: #4CAF50;
|
||||
color: white;
|
||||
padding: 20px;
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
.connection-status {
|
||||
padding: 10px;
|
||||
text-align: center;
|
||||
font-weight: bold;
|
||||
}
|
||||
|
||||
.connected {
|
||||
background-color: #d4edda;
|
||||
color: #155724;
|
||||
}
|
||||
|
||||
.disconnected {
|
||||
background-color: #f8d7da;
|
||||
color: #721c24;
|
||||
}
|
||||
|
||||
.connecting {
|
||||
background-color: #fff3cd;
|
||||
color: #856404;
|
||||
}
|
||||
|
||||
.messages {
|
||||
height: 400px;
|
||||
overflow-y: auto;
|
||||
padding: 20px;
|
||||
border-bottom: 1px solid #eee;
|
||||
}
|
||||
|
||||
.message {
|
||||
margin-bottom: 15px;
|
||||
padding: 10px;
|
||||
border-radius: 5px;
|
||||
word-wrap: break-word;
|
||||
}
|
||||
|
||||
.message.chat {
|
||||
background-color: #e3f2fd;
|
||||
border-left: 4px solid #2196F3;
|
||||
}
|
||||
|
||||
.message.system {
|
||||
background-color: #f3e5f5;
|
||||
border-left: 4px solid #9c27b0;
|
||||
font-style: italic;
|
||||
}
|
||||
|
||||
.message.error {
|
||||
background-color: #ffebee;
|
||||
border-left: 4px solid #f44336;
|
||||
}
|
||||
|
||||
.message-sender {
|
||||
font-weight: bold;
|
||||
color: #333;
|
||||
margin-bottom: 5px;
|
||||
}
|
||||
|
||||
.message-content {
|
||||
color: #666;
|
||||
}
|
||||
|
||||
.message-timestamp {
|
||||
font-size: 0.8em;
|
||||
color: #999;
|
||||
margin-top: 5px;
|
||||
}
|
||||
|
||||
.input-area {
|
||||
padding: 20px;
|
||||
background-color: #fafafa;
|
||||
}
|
||||
|
||||
.input-row {
|
||||
display: flex;
|
||||
gap: 10px;
|
||||
margin-bottom: 10px;
|
||||
}
|
||||
|
||||
#usernameInput {
|
||||
flex: 1;
|
||||
padding: 10px;
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 5px;
|
||||
font-size: 14px;
|
||||
}
|
||||
|
||||
#messageInput {
|
||||
flex: 3;
|
||||
padding: 10px;
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 5px;
|
||||
font-size: 14px;
|
||||
}
|
||||
|
||||
#sendButton {
|
||||
padding: 10px 20px;
|
||||
background: #4CAF50;
|
||||
color: white;
|
||||
border: none;
|
||||
border-radius: 5px;
|
||||
cursor: pointer;
|
||||
font-size: 14px;
|
||||
}
|
||||
|
||||
#sendButton:hover:not(:disabled) {
|
||||
background: #45a049;
|
||||
}
|
||||
|
||||
#sendButton:disabled {
|
||||
background: #cccccc;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
.controls {
|
||||
display: flex;
|
||||
gap: 10px;
|
||||
justify-content: center;
|
||||
}
|
||||
|
||||
.btn {
|
||||
padding: 8px 16px;
|
||||
border: none;
|
||||
border-radius: 5px;
|
||||
cursor: pointer;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.btn-primary {
|
||||
background: #2196F3;
|
||||
color: white;
|
||||
}
|
||||
|
||||
.btn-primary:hover {
|
||||
background: #1976D2;
|
||||
}
|
||||
|
||||
.btn-danger {
|
||||
background: #f44336;
|
||||
color: white;
|
||||
}
|
||||
|
||||
.btn-danger:hover {
|
||||
background: #d32f2f;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="chat-container">
|
||||
<div class="chat-header">
|
||||
<h1>WebSocket Chat Room</h1>
|
||||
</div>
|
||||
|
||||
<div id="connectionStatus" class="connection-status disconnected">
|
||||
Disconnected
|
||||
</div>
|
||||
|
||||
<div id="messages" class="messages">
|
||||
<div class="message system">
|
||||
<div class="message-content">Welcome to the chat room! Enter your username and click Connect to start chatting.</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="input-area">
|
||||
<div class="input-row">
|
||||
<input type="text" id="usernameInput" placeholder="Enter your username..." value="Anonymous">
|
||||
<input type="text" id="messageInput" placeholder="Type your message..." disabled>
|
||||
<button id="sendButton" disabled>Send</button>
|
||||
</div>
|
||||
<div class="controls">
|
||||
<button id="connectButton" class="btn btn-primary">Connect</button>
|
||||
<button id="disconnectButton" class="btn btn-danger" disabled>Disconnect</button>
|
||||
<button id="clearButton" class="btn">Clear Messages</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
class ChatClient {
|
||||
constructor() {
|
||||
this.websocket = null;
|
||||
this.isConnected = false;
|
||||
this.username = 'Anonymous';
|
||||
this.serverUrl = 'ws://localhost:8765';
|
||||
|
||||
this.initializeElements();
|
||||
this.attachEventListeners();
|
||||
}
|
||||
|
||||
initializeElements() {
|
||||
this.elements = {
|
||||
connectionStatus: document.getElementById('connectionStatus'),
|
||||
messages: document.getElementById('messages'),
|
||||
usernameInput: document.getElementById('usernameInput'),
|
||||
messageInput: document.getElementById('messageInput'),
|
||||
sendButton: document.getElementById('sendButton'),
|
||||
connectButton: document.getElementById('connectButton'),
|
||||
disconnectButton: document.getElementById('disconnectButton'),
|
||||
clearButton: document.getElementById('clearButton')
|
||||
};
|
||||
}
|
||||
|
||||
attachEventListeners() {
|
||||
this.elements.connectButton.addEventListener('click', () => this.connect());
|
||||
this.elements.disconnectButton.addEventListener('click', () => this.disconnect());
|
||||
this.elements.sendButton.addEventListener('click', () => this.sendMessage());
|
||||
this.elements.clearButton.addEventListener('click', () => this.clearMessages());
|
||||
|
||||
this.elements.messageInput.addEventListener('keypress', (e) => {
|
||||
if (e.key === 'Enter' && !e.shiftKey) {
|
||||
e.preventDefault();
|
||||
this.sendMessage();
|
||||
}
|
||||
});
|
||||
|
||||
this.elements.usernameInput.addEventListener('keypress', (e) => {
|
||||
if (e.key === 'Enter') {
|
||||
this.connect();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async connect() {
|
||||
if (this.isConnected) return;
|
||||
|
||||
this.username = this.elements.usernameInput.value.trim() || 'Anonymous';
|
||||
this.updateConnectionStatus('connecting', 'Connecting...');
|
||||
|
||||
try {
|
||||
this.websocket = new WebSocket(this.serverUrl);
|
||||
|
||||
this.websocket.onopen = () => {
|
||||
this.isConnected = true;
|
||||
this.updateConnectionStatus('connected', 'Connected');
|
||||
this.updateButtonStates();
|
||||
this.addSystemMessage('Connected to chat server');
|
||||
};
|
||||
|
||||
this.websocket.onmessage = (event) => {
|
||||
this.handleMessage(event.data);
|
||||
};
|
||||
|
||||
this.websocket.onclose = (event) => {
|
||||
this.isConnected = false;
|
||||
this.updateConnectionStatus('disconnected', 'Disconnected');
|
||||
this.updateButtonStates();
|
||||
|
||||
if (event.wasClean) {
|
||||
this.addSystemMessage('Disconnected from chat server');
|
||||
} else {
|
||||
this.addSystemMessage('Connection lost unexpectedly');
|
||||
}
|
||||
|
||||
this.websocket = null;
|
||||
};
|
||||
|
||||
this.websocket.onerror = (error) => {
|
||||
console.error('WebSocket error:', error);
|
||||
this.addErrorMessage('Connection error occurred');
|
||||
};
|
||||
|
||||
} catch (error) {
|
||||
console.error('Failed to connect:', error);
|
||||
this.updateConnectionStatus('disconnected', 'Connection failed');
|
||||
this.addErrorMessage('Failed to connect to server');
|
||||
}
|
||||
}
|
||||
|
||||
disconnect() {
|
||||
if (this.websocket && this.isConnected) {
|
||||
this.websocket.close(1000, 'User disconnected');
|
||||
}
|
||||
}
|
||||
|
||||
sendMessage() {
|
||||
const message = this.elements.messageInput.value.trim();
|
||||
if (!message || !this.isConnected) return;
|
||||
|
||||
try {
|
||||
const messageData = {
|
||||
message: message,
|
||||
sender: this.username
|
||||
};
|
||||
|
||||
this.websocket.send(JSON.stringify(messageData));
|
||||
this.elements.messageInput.value = '';
|
||||
|
||||
} catch (error) {
|
||||
console.error('Failed to send message:', error);
|
||||
this.addErrorMessage('Failed to send message');
|
||||
}
|
||||
}
|
||||
|
||||
handleMessage(data) {
|
||||
try {
|
||||
const message = JSON.parse(data);
|
||||
|
||||
switch (message.type) {
|
||||
case 'chat':
|
||||
this.addChatMessage(message);
|
||||
break;
|
||||
case 'system':
|
||||
this.addSystemMessage(message.message);
|
||||
break;
|
||||
case 'error':
|
||||
this.addErrorMessage(message.message);
|
||||
break;
|
||||
default:
|
||||
console.warn('Unknown message type:', message.type);
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Failed to parse message:', error);
|
||||
this.addErrorMessage('Received invalid message format');
|
||||
}
|
||||
}
|
||||
|
||||
addChatMessage(message) {
|
||||
const messageEl = this.createMessageElement('chat');
|
||||
messageEl.innerHTML = `
|
||||
<div class="message-sender">${this.escapeHtml(message.sender)}</div>
|
||||
<div class="message-content">${this.escapeHtml(message.message)}</div>
|
||||
<div class="message-timestamp">${this.formatTimestamp(message.timestamp)}</div>
|
||||
`;
|
||||
this.elements.messages.appendChild(messageEl);
|
||||
this.scrollToBottom();
|
||||
}
|
||||
|
||||
addSystemMessage(text) {
|
||||
const messageEl = this.createMessageElement('system');
|
||||
messageEl.innerHTML = `
|
||||
<div class="message-content">${this.escapeHtml(text)}</div>
|
||||
<div class="message-timestamp">${this.formatTimestamp(Date.now() / 1000)}</div>
|
||||
`;
|
||||
this.elements.messages.appendChild(messageEl);
|
||||
this.scrollToBottom();
|
||||
}
|
||||
|
||||
addErrorMessage(text) {
|
||||
const messageEl = this.createMessageElement('error');
|
||||
messageEl.innerHTML = `
|
||||
<div class="message-content">Error: ${this.escapeHtml(text)}</div>
|
||||
<div class="message-timestamp">${this.formatTimestamp(Date.now() / 1000)}</div>
|
||||
`;
|
||||
this.elements.messages.appendChild(messageEl);
|
||||
this.scrollToBottom();
|
||||
}
|
||||
|
||||
createMessageElement(type) {
|
||||
const messageEl = document.createElement('div');
|
||||
messageEl.className = `message ${type}`;
|
||||
return messageEl;
|
||||
}
|
||||
|
||||
updateConnectionStatus(status, text) {
|
||||
this.elements.connectionStatus.className = `connection-status ${status}`;
|
||||
this.elements.connectionStatus.textContent = text;
|
||||
}
|
||||
|
||||
updateButtonStates() {
|
||||
this.elements.connectButton.disabled = this.isConnected;
|
||||
this.elements.disconnectButton.disabled = !this.isConnected;
|
||||
this.elements.sendButton.disabled = !this.isConnected;
|
||||
this.elements.messageInput.disabled = !this.isConnected;
|
||||
this.elements.usernameInput.disabled = this.isConnected;
|
||||
}
|
||||
|
||||
clearMessages() {
|
||||
this.elements.messages.innerHTML = '';
|
||||
this.addSystemMessage('Messages cleared');
|
||||
}
|
||||
|
||||
scrollToBottom() {
|
||||
this.elements.messages.scrollTop = this.elements.messages.scrollHeight;
|
||||
}
|
||||
|
||||
escapeHtml(text) {
|
||||
const div = document.createElement('div');
|
||||
div.textContent = text;
|
||||
return div.innerHTML;
|
||||
}
|
||||
|
||||
formatTimestamp(timestamp) {
|
||||
const date = new Date(timestamp * 1000);
|
||||
return date.toLocaleTimeString();
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize the chat client when the page loads
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
window.chatClient = new ChatClient();
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1 @@
|
||||
websockets>=11.0.0
|
||||
@@ -0,0 +1,139 @@
|
||||
import asyncio
|
||||
import websockets
|
||||
import json
|
||||
import logging
|
||||
from typing import Set
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global set to store connected clients
|
||||
connected_clients: Set[websockets.WebSocketServerProtocol] = set()
|
||||
|
||||
async def register_client(websocket: websockets.WebSocketServerProtocol) -> None:
|
||||
"""Register a new client connection."""
|
||||
connected_clients.add(websocket)
|
||||
logger.info(f"Client {websocket.remote_address} connected. Total clients: {len(connected_clients)}")
|
||||
|
||||
# Notify all clients about new connection
|
||||
if len(connected_clients) > 1:
|
||||
notification = {
|
||||
"type": "system",
|
||||
"message": f"User from {websocket.remote_address[0]} joined the chat",
|
||||
"timestamp": asyncio.get_event_loop().time()
|
||||
}
|
||||
await broadcast_message(json.dumps(notification), exclude=websocket)
|
||||
|
||||
async def unregister_client(websocket: websockets.WebSocketServerProtocol) -> None:
|
||||
"""Unregister a client connection."""
|
||||
if websocket in connected_clients:
|
||||
connected_clients.remove(websocket)
|
||||
logger.info(f"Client {websocket.remote_address} disconnected. Total clients: {len(connected_clients)}")
|
||||
|
||||
# Notify remaining clients about disconnection
|
||||
if connected_clients:
|
||||
notification = {
|
||||
"type": "system",
|
||||
"message": f"User from {websocket.remote_address[0]} left the chat",
|
||||
"timestamp": asyncio.get_event_loop().time()
|
||||
}
|
||||
await broadcast_message(json.dumps(notification))
|
||||
|
||||
async def broadcast_message(message: str, exclude: websockets.WebSocketServerProtocol = None) -> None:
|
||||
"""Broadcast a message to all connected clients except the excluded one."""
|
||||
if not connected_clients:
|
||||
return
|
||||
|
||||
# Create a copy of the set to avoid modification during iteration
|
||||
clients_to_notify = connected_clients.copy()
|
||||
if exclude:
|
||||
clients_to_notify.discard(exclude)
|
||||
|
||||
if clients_to_notify:
|
||||
# Send message to all clients concurrently
|
||||
await asyncio.gather(
|
||||
*[send_safe(client, message) for client in clients_to_notify],
|
||||
return_exceptions=True
|
||||
)
|
||||
|
||||
async def send_safe(websocket: websockets.WebSocketServerProtocol, message: str) -> None:
|
||||
"""Safely send a message to a client with error handling."""
|
||||
try:
|
||||
await websocket.send(message)
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
logger.warning(f"Failed to send message to {websocket.remote_address}: connection closed")
|
||||
await unregister_client(websocket)
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending message to {websocket.remote_address}: {e}")
|
||||
await unregister_client(websocket)
|
||||
|
||||
async def handle_client(websocket: websockets.WebSocketServerProtocol, path: str) -> None:
|
||||
"""Handle a client connection."""
|
||||
await register_client(websocket)
|
||||
|
||||
try:
|
||||
async for message in websocket:
|
||||
try:
|
||||
# Parse the incoming message
|
||||
data = json.loads(message)
|
||||
|
||||
# Create a broadcast message with metadata
|
||||
broadcast_data = {
|
||||
"type": "chat",
|
||||
"message": data.get("message", ""),
|
||||
"sender": data.get("sender", "Anonymous"),
|
||||
"timestamp": asyncio.get_event_loop().time()
|
||||
}
|
||||
|
||||
logger.info(f"Broadcasting message from {broadcast_data['sender']}: {broadcast_data['message']}")
|
||||
|
||||
# Broadcast to all connected clients
|
||||
await broadcast_message(json.dumps(broadcast_data))
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Invalid JSON received from {websocket.remote_address}")
|
||||
error_response = {
|
||||
"type": "error",
|
||||
"message": "Invalid message format",
|
||||
"timestamp": asyncio.get_event_loop().time()
|
||||
}
|
||||
await send_safe(websocket, json.dumps(error_response))
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing message from {websocket.remote_address}: {e}")
|
||||
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
logger.info(f"Connection closed by {websocket.remote_address}")
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error with {websocket.remote_address}: {e}")
|
||||
finally:
|
||||
await unregister_client(websocket)
|
||||
|
||||
async def main():
|
||||
"""Start the WebSocket server."""
|
||||
host = "localhost"
|
||||
port = 8765
|
||||
|
||||
logger.info(f"Starting WebSocket server on {host}:{port}")
|
||||
|
||||
try:
|
||||
async with websockets.serve(handle_client, host, port):
|
||||
logger.info(f"WebSocket server started successfully on ws://{host}:{port}")
|
||||
logger.info("Press Ctrl+C to stop the server")
|
||||
|
||||
# Keep the server running
|
||||
await asyncio.Future() # Run forever
|
||||
|
||||
except OSError as e:
|
||||
logger.error(f"Failed to start server: {e}")
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Server shutdown requested")
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected server error: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Server stopped")
|
||||
@@ -0,0 +1,320 @@
|
||||
I'll create a production-ready Hacker News scraper that extracts the top 10 stories and saves them to JSON format.
|
||||
|
||||
## hacker_news_scraper.py
|
||||
|
||||
```python
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Hacker News Top Stories Scraper
|
||||
|
||||
Scrapes the top 10 stories from Hacker News front page and saves them to JSON.
|
||||
Includes rate limiting, error handling, and robust data extraction.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Dict, List, Optional
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
|
||||
class HackerNewsScraper:
|
||||
"""Scraper for Hacker News top stories."""
|
||||
|
||||
def __init__(self, delay: float = 1.0, timeout: int = 10):
|
||||
"""
|
||||
Initialize the scraper.
|
||||
|
||||
Args:
|
||||
delay: Delay between requests in seconds
|
||||
timeout: Request timeout in seconds
|
||||
"""
|
||||
self.base_url = "https://news.ycombinator.com"
|
||||
self.delay = delay
|
||||
self.timeout = timeout
|
||||
self.session = requests.Session()
|
||||
|
||||
# Set user agent to be respectful
|
||||
self.session.headers.update({
|
||||
'User-Agent': 'HackerNewsScraper/1.0 (Educational Purpose)'
|
||||
})
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
def _make_request(self, url: str) -> Optional[requests.Response]:
|
||||
"""
|
||||
Make HTTP request with error handling and rate limiting.
|
||||
|
||||
Args:
|
||||
url: URL to request
|
||||
|
||||
Returns:
|
||||
Response object or None if request failed
|
||||
"""
|
||||
try:
|
||||
self.logger.info(f"Requesting: {url}")
|
||||
response = self.session.get(url, timeout=self.timeout)
|
||||
response.raise_for_status()
|
||||
|
||||
# Rate limiting
|
||||
time.sleep(self.delay)
|
||||
|
||||
return response
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
self.logger.error(f"Request failed for {url}: {e}")
|
||||
return None
|
||||
|
||||
def _extract_story_data(self, story_row, subtext_row) -> Optional[Dict]:
|
||||
"""
|
||||
Extract data from story and subtext rows.
|
||||
|
||||
Args:
|
||||
story_row: BeautifulSoup element containing story title/link
|
||||
subtext_row: BeautifulSoup element containing metadata
|
||||
|
||||
Returns:
|
||||
Dictionary with story data or None if extraction failed
|
||||
"""
|
||||
try:
|
||||
# Extract title and URL
|
||||
title_elem = story_row.select_one('.titleline > a')
|
||||
if not title_elem:
|
||||
return None
|
||||
|
||||
title = title_elem.get_text(strip=True)
|
||||
url = title_elem.get('href', '')
|
||||
|
||||
# Handle relative URLs
|
||||
if url.startswith('item?'):
|
||||
url = urljoin(self.base_url, url)
|
||||
elif not url.startswith('http'):
|
||||
url = urljoin(self.base_url, url)
|
||||
|
||||
# Extract domain
|
||||
domain_elem = story_row.select_one('.sitestr')
|
||||
domain = domain_elem.get_text(strip=True) if domain_elem else ''
|
||||
|
||||
# Extract metadata from subtext
|
||||
score_elem = subtext_row.select_one('.score')
|
||||
score = 0
|
||||
if score_elem:
|
||||
score_text = score_elem.get_text(strip=True)
|
||||
score = int(score_text.split()[0]) if score_text else 0
|
||||
|
||||
# Extract author
|
||||
author_elem = subtext_row.select_one('.hnuser')
|
||||
author = author_elem.get_text(strip=True) if author_elem else ''
|
||||
|
||||
# Extract comments count
|
||||
comments_elem = subtext_row.select('a')[-1] # Last link is usually comments
|
||||
comments_text = comments_elem.get_text(strip=True) if comments_elem else ''
|
||||
comments = 0
|
||||
|
||||
if 'comment' in comments_text:
|
||||
try:
|
||||
comments = int(comments_text.split()[0])
|
||||
except (ValueError, IndexError):
|
||||
comments = 0
|
||||
|
||||
# Extract time
|
||||
time_elem = subtext_row.select_one('.age')
|
||||
time_posted = time_elem.get('title', '') if time_elem else ''
|
||||
|
||||
return {
|
||||
'title': title,
|
||||
'url': url,
|
||||
'domain': domain,
|
||||
'score': score,
|
||||
'author': author,
|
||||
'comments': comments,
|
||||
'time_posted': time_posted
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error extracting story data: {e}")
|
||||
return None
|
||||
|
||||
def scrape_top_stories(self, count: int = 10) -> List[Dict]:
|
||||
"""
|
||||
Scrape top stories from Hacker News.
|
||||
|
||||
Args:
|
||||
count: Number of top stories to retrieve
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing story data
|
||||
"""
|
||||
response = self._make_request(self.base_url)
|
||||
if not response:
|
||||
self.logger.error("Failed to fetch Hacker News homepage")
|
||||
return []
|
||||
|
||||
soup = BeautifulSoup(response.text, 'html.parser')
|
||||
stories = []
|
||||
|
||||
# Find story rows - they have class 'athing'
|
||||
story_rows = soup.select('.athing')[:count]
|
||||
|
||||
for i, story_row in enumerate(story_rows):
|
||||
try:
|
||||
# Get the corresponding subtext row (contains score, author, etc.)
|
||||
story_id = story_row.get('id')
|
||||
if not story_id:
|
||||
continue
|
||||
|
||||
# Find the subtext row (next sibling with story metadata)
|
||||
subtext_row = story_row.find_next_sibling('tr')
|
||||
if not subtext_row:
|
||||
continue
|
||||
|
||||
story_data = self._extract_story_data(story_row, subtext_row)
|
||||
if story_data:
|
||||
story_data['rank'] = i + 1
|
||||
stories.append(story_data)
|
||||
self.logger.info(f"Extracted story {i+1}: {story_data['title']}")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing story {i+1}: {e}")
|
||||
continue
|
||||
|
||||
self.logger.info(f"Successfully scraped {len(stories)} stories")
|
||||
return stories
|
||||
|
||||
def save_to_json(self, stories: List[Dict], filename: str = 'stories.json') -> bool:
|
||||
"""
|
||||
Save stories to JSON file.
|
||||
|
||||
Args:
|
||||
stories: List of story dictionaries
|
||||
filename: Output filename
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
output_data = {
|
||||
'scraped_at': time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime()),
|
||||
'total_stories': len(stories),
|
||||
'stories': stories
|
||||
}
|
||||
|
||||
with open(filename, 'w', encoding='utf-8') as f:
|
||||
json.dump(output_data, f, indent=2, ensure_ascii=False)
|
||||
|
||||
self.logger.info(f"Stories saved to {filename}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error saving to JSON: {e}")
|
||||
return False
|
||||
|
||||
def close(self):
|
||||
"""Close the session."""
|
||||
self.session.close()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to run the scraper."""
|
||||
scraper = HackerNewsScraper(delay=1.0)
|
||||
|
||||
try:
|
||||
# Scrape top 10 stories
|
||||
stories = scraper.scrape_top_stories(count=10)
|
||||
|
||||
if stories:
|
||||
# Save to JSON
|
||||
if scraper.save_to_json(stories):
|
||||
print(f"Successfully scraped {len(stories)} stories and saved to stories.json")
|
||||
|
||||
# Print summary
|
||||
print("\nTop Stories Summary:")
|
||||
print("-" * 50)
|
||||
for story in stories:
|
||||
print(f"{story['rank']}. {story['title']}")
|
||||
print(f" Score: {story['score']} | Comments: {story['comments']} | Author: {story['author']}")
|
||||
print()
|
||||
else:
|
||||
print("Failed to save stories to JSON file")
|
||||
else:
|
||||
print("No stories were scraped")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nScraping interrupted by user")
|
||||
except Exception as e:
|
||||
print(f"Unexpected error: {e}")
|
||||
finally:
|
||||
scraper.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
## requirements.txt
|
||||
|
||||
```txt
|
||||
requests>=2.31.0
|
||||
beautifulsoup4>=4.12.0
|
||||
lxml>=4.9.0
|
||||
```
|
||||
|
||||
## Usage Example
|
||||
|
||||
```python
|
||||
# Example of using the scraper programmatically
|
||||
from hacker_news_scraper import HackerNewsScraper
|
||||
|
||||
# Initialize scraper
|
||||
scraper = HackerNewsScraper(delay=1.5) # 1.5 second delay between requests
|
||||
|
||||
# Scrape top 10 stories
|
||||
stories = scraper.scrape_top_stories(count=10)
|
||||
|
||||
# Save to custom filename
|
||||
scraper.save_to_json(stories, 'hn_top_stories.json')
|
||||
|
||||
# Clean up
|
||||
scraper.close()
|
||||
```
|
||||
|
||||
## Key Features
|
||||
|
||||
1. **Rate Limiting**: 1-second delay between requests to be respectful
|
||||
2. **Error Handling**: Comprehensive error handling for network issues and parsing errors
|
||||
3. **Robust Data Extraction**: Handles missing fields and malformed data gracefully
|
||||
4. **Logging**: Detailed logging for debugging and monitoring
|
||||
5. **Session Management**: Reuses HTTP connections for efficiency
|
||||
6. **Configurable**: Easy to adjust timeouts, delays, and story count
|
||||
7. **Production Ready**: Includes proper documentation, type hints, and error recovery
|
||||
|
||||
## Output Format (stories.json)
|
||||
|
||||
```json
|
||||
{
|
||||
"scraped_at": "2024-01-15 14:30:25 UTC",
|
||||
"total_stories": 10,
|
||||
"stories": [
|
||||
{
|
||||
"rank": 1,
|
||||
"title": "Example Story Title",
|
||||
"url": "https://example.com/article",
|
||||
"domain": "example.com",
|
||||
"score": 150,
|
||||
"author": "username",
|
||||
"comments": 45,
|
||||
"time_posted": "2024-01-15T13:25:00"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
The scraper is respectful of Hacker News's servers, includes proper error handling, and produces clean, structured JSON output suitable for further processing.
|
||||
@@ -0,0 +1,3 @@
|
||||
requests>=2.28.0
|
||||
beautifulsoup4>=4.11.0
|
||||
lxml>=4.9.0
|
||||
@@ -0,0 +1,251 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Hacker News Scraper
|
||||
|
||||
A web scraper that fetches stories from the Hacker News front page
|
||||
and outputs them in JSON format.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass, asdict
|
||||
from typing import List, Optional
|
||||
from urllib.parse import urljoin, urlparse
|
||||
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Constants
|
||||
HACKER_NEWS_URL = "https://news.ycombinator.com/"
|
||||
REQUEST_TIMEOUT = 10
|
||||
USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Story:
|
||||
"""Represents a Hacker News story with metadata."""
|
||||
title: str
|
||||
url: str
|
||||
points: int
|
||||
comments: int
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert Story to dictionary for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
class HackerNewsScraper:
|
||||
"""Scraper for Hacker News front page stories."""
|
||||
|
||||
def __init__(self):
|
||||
self.session = requests.Session()
|
||||
self.session.headers.update({'User-Agent': USER_AGENT})
|
||||
|
||||
def fetch_page(self, url: str) -> Optional[str]:
|
||||
"""
|
||||
Fetch HTML content from the given URL.
|
||||
|
||||
Args:
|
||||
url: The URL to fetch
|
||||
|
||||
Returns:
|
||||
HTML content as string, or None if request failed
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Fetching page: {url}")
|
||||
response = self.session.get(url, timeout=REQUEST_TIMEOUT)
|
||||
response.raise_for_status()
|
||||
return response.text
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Failed to fetch {url}: {e}")
|
||||
return None
|
||||
|
||||
def extract_number(self, text: str) -> int:
|
||||
"""
|
||||
Extract numeric value from text string.
|
||||
|
||||
Args:
|
||||
text: Text containing a number
|
||||
|
||||
Returns:
|
||||
Extracted number, or 0 if no number found
|
||||
"""
|
||||
if not text:
|
||||
return 0
|
||||
|
||||
match = re.search(r'(\d+)', text.strip())
|
||||
return int(match.group(1)) if match else 0
|
||||
|
||||
def parse_story_row(self, title_row, subtext_row) -> Optional[Story]:
|
||||
"""
|
||||
Parse a story from title row and subtext row elements.
|
||||
|
||||
Args:
|
||||
title_row: BeautifulSoup element containing title and URL
|
||||
subtext_row: BeautifulSoup element containing points and comments
|
||||
|
||||
Returns:
|
||||
Story object or None if parsing failed
|
||||
"""
|
||||
try:
|
||||
# Extract title and URL
|
||||
title_link = title_row.find('span', class_='titleline')
|
||||
if not title_link:
|
||||
return None
|
||||
|
||||
link_elem = title_link.find('a')
|
||||
if not link_elem:
|
||||
return None
|
||||
|
||||
title = link_elem.get_text(strip=True)
|
||||
url = link_elem.get('href', '')
|
||||
|
||||
# Handle relative URLs
|
||||
if url.startswith('item?'):
|
||||
url = urljoin(HACKER_NEWS_URL, url)
|
||||
elif not urlparse(url).netloc:
|
||||
url = urljoin(HACKER_NEWS_URL, url)
|
||||
|
||||
# Extract points
|
||||
points = 0
|
||||
score_elem = subtext_row.find('span', class_='score')
|
||||
if score_elem:
|
||||
points = self.extract_number(score_elem.get_text())
|
||||
|
||||
# Extract comments count
|
||||
comments = 0
|
||||
comment_links = subtext_row.find_all('a')
|
||||
for link in comment_links:
|
||||
link_text = link.get_text(strip=True)
|
||||
if 'comment' in link_text.lower():
|
||||
comments = self.extract_number(link_text)
|
||||
break
|
||||
|
||||
return Story(
|
||||
title=title,
|
||||
url=url,
|
||||
points=points,
|
||||
comments=comments
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to parse story row: {e}")
|
||||
return None
|
||||
|
||||
def parse_stories(self, html: str) -> List[Story]:
|
||||
"""
|
||||
Parse stories from Hacker News HTML content.
|
||||
|
||||
Args:
|
||||
html: HTML content string
|
||||
|
||||
Returns:
|
||||
List of Story objects
|
||||
"""
|
||||
try:
|
||||
soup = BeautifulSoup(html, 'lxml')
|
||||
stories = []
|
||||
|
||||
# Find the main table containing stories
|
||||
main_table = soup.find('table', id='hnmain')
|
||||
if not main_table:
|
||||
logger.error("Could not find main stories table")
|
||||
return stories
|
||||
|
||||
# Find all story rows (they have class 'athing')
|
||||
story_rows = main_table.find_all('tr', class_='athing')
|
||||
|
||||
for story_row in story_rows:
|
||||
# The subtext row immediately follows the story row
|
||||
subtext_row = story_row.find_next_sibling('tr')
|
||||
|
||||
if subtext_row and subtext_row.find('td', class_='subtext'):
|
||||
story = self.parse_story_row(story_row, subtext_row)
|
||||
if story:
|
||||
stories.append(story)
|
||||
logger.debug(f"Parsed story: {story.title}")
|
||||
|
||||
logger.info(f"Successfully parsed {len(stories)} stories")
|
||||
return stories
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse HTML: {e}")
|
||||
return []
|
||||
|
||||
def scrape_front_page(self) -> List[Story]:
|
||||
"""
|
||||
Scrape stories from Hacker News front page.
|
||||
|
||||
Returns:
|
||||
List of Story objects
|
||||
"""
|
||||
html = self.fetch_page(HACKER_NEWS_URL)
|
||||
if not html:
|
||||
logger.error("Failed to fetch Hacker News front page")
|
||||
return []
|
||||
|
||||
return self.parse_stories(html)
|
||||
|
||||
def close(self):
|
||||
"""Close the requests session."""
|
||||
self.session.close()
|
||||
|
||||
|
||||
def stories_to_json(stories: List[Story]) -> str:
|
||||
"""
|
||||
Convert list of stories to JSON string.
|
||||
|
||||
Args:
|
||||
stories: List of Story objects
|
||||
|
||||
Returns:
|
||||
JSON string representation
|
||||
"""
|
||||
try:
|
||||
stories_dict = [story.to_dict() for story in stories]
|
||||
return json.dumps(stories_dict, indent=2, ensure_ascii=False)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to serialize stories to JSON: {e}")
|
||||
return "[]"
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to run the scraper."""
|
||||
scraper = HackerNewsScraper()
|
||||
|
||||
try:
|
||||
# Scrape stories
|
||||
stories = scraper.scrape_front_page()
|
||||
|
||||
if not stories:
|
||||
logger.warning("No stories were scraped")
|
||||
print("[]")
|
||||
sys.exit(1)
|
||||
|
||||
# Output as JSON
|
||||
json_output = stories_to_json(stories)
|
||||
print(json_output)
|
||||
|
||||
logger.info(f"Successfully scraped and output {len(stories)} stories")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Scraping interrupted by user")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error: {e}")
|
||||
sys.exit(1)
|
||||
finally:
|
||||
scraper.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,552 @@
|
||||
I'll create a production-ready Flask web application with user authentication. Here's the complete implementation:
|
||||
|
||||
## requirements.txt
|
||||
```txt
|
||||
Flask>=2.3.0
|
||||
Flask-Login>=0.6.0
|
||||
Flask-SQLAlchemy>=3.0.0
|
||||
Flask-WTF>=1.1.0
|
||||
WTForms>=3.0.0
|
||||
Werkzeug>=2.3.0
|
||||
bcrypt>=4.0.0
|
||||
email-validator>=2.0.0
|
||||
```
|
||||
|
||||
## config.py
|
||||
```python
|
||||
import os
|
||||
from datetime import timedelta
|
||||
|
||||
class Config:
|
||||
"""Application configuration class"""
|
||||
SECRET_KEY = os.environ.get('SECRET_KEY') or 'dev-key-change-in-production'
|
||||
SQLALCHEMY_DATABASE_URI = os.environ.get('DATABASE_URL') or 'sqlite:///app.db'
|
||||
SQLALCHEMY_TRACK_MODIFICATIONS = False
|
||||
|
||||
# Session configuration
|
||||
PERMANENT_SESSION_LIFETIME = timedelta(hours=24)
|
||||
SESSION_COOKIE_SECURE = os.environ.get('FLASK_ENV') == 'production'
|
||||
SESSION_COOKIE_HTTPONLY = True
|
||||
SESSION_COOKIE_SAMESITE = 'Lax'
|
||||
|
||||
# WTF-Forms CSRF protection
|
||||
WTF_CSRF_ENABLED = True
|
||||
WTF_CSRF_TIME_LIMIT = 3600
|
||||
```
|
||||
|
||||
## models.py
|
||||
```python
|
||||
from flask_sqlalchemy import SQLAlchemy
|
||||
from flask_login import UserMixin
|
||||
from werkzeug.security import generate_password_hash, check_password_hash
|
||||
from datetime import datetime
|
||||
|
||||
db = SQLAlchemy()
|
||||
|
||||
class User(UserMixin, db.Model):
|
||||
"""User model for authentication"""
|
||||
__tablename__ = 'users'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
username = db.Column(db.String(80), unique=True, nullable=False, index=True)
|
||||
email = db.Column(db.String(120), unique=True, nullable=False, index=True)
|
||||
password_hash = db.Column(db.String(255), nullable=False)
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
is_active = db.Column(db.Boolean, default=True)
|
||||
|
||||
def set_password(self, password):
|
||||
"""Hash and set password"""
|
||||
self.password_hash = generate_password_hash(password, method='pbkdf2:sha256', salt_length=16)
|
||||
|
||||
def check_password(self, password):
|
||||
"""Check if provided password matches hash"""
|
||||
return check_password_hash(self.password_hash, password)
|
||||
|
||||
def get_id(self):
|
||||
"""Return user id as string for Flask-Login"""
|
||||
return str(self.id)
|
||||
|
||||
def __repr__(self):
|
||||
return f'<User {self.username}>'
|
||||
```
|
||||
|
||||
## forms.py
|
||||
```python
|
||||
from flask_wtf import FlaskForm
|
||||
from wtforms import StringField, PasswordField, SubmitField, BooleanField
|
||||
from wtforms.validators import DataRequired, Email, Length, EqualTo, ValidationError
|
||||
from models import User
|
||||
|
||||
class LoginForm(FlaskForm):
|
||||
"""User login form"""
|
||||
username = StringField('Username', validators=[
|
||||
DataRequired(message='Username is required'),
|
||||
Length(min=3, max=20, message='Username must be between 3 and 20 characters')
|
||||
])
|
||||
password = PasswordField('Password', validators=[
|
||||
DataRequired(message='Password is required')
|
||||
])
|
||||
remember_me = BooleanField('Remember Me')
|
||||
submit = SubmitField('Sign In')
|
||||
|
||||
class RegistrationForm(FlaskForm):
|
||||
"""User registration form"""
|
||||
username = StringField('Username', validators=[
|
||||
DataRequired(message='Username is required'),
|
||||
Length(min=3, max=20, message='Username must be between 3 and 20 characters')
|
||||
])
|
||||
email = StringField('Email', validators=[
|
||||
DataRequired(message='Email is required'),
|
||||
Email(message='Please enter a valid email address')
|
||||
])
|
||||
password = PasswordField('Password', validators=[
|
||||
DataRequired(message='Password is required'),
|
||||
Length(min=8, message='Password must be at least 8 characters long')
|
||||
])
|
||||
password2 = PasswordField('Repeat Password', validators=[
|
||||
DataRequired(message='Please confirm your password'),
|
||||
EqualTo('password', message='Passwords must match')
|
||||
])
|
||||
submit = SubmitField('Register')
|
||||
|
||||
def validate_username(self, username):
|
||||
"""Check if username already exists"""
|
||||
user = User.query.filter_by(username=username.data).first()
|
||||
if user:
|
||||
raise ValidationError('Username already exists. Please choose a different one.')
|
||||
|
||||
def validate_email(self, email):
|
||||
"""Check if email already exists"""
|
||||
user = User.query.filter_by(email=email.data).first()
|
||||
if user:
|
||||
raise ValidationError('Email already registered. Please use a different email.')
|
||||
```
|
||||
|
||||
## app.py
|
||||
```python
|
||||
from flask import Flask, render_template, redirect, url_for, flash, request
|
||||
from flask_login import LoginManager, login_user, logout_user, login_required, current_user
|
||||
from models import db, User
|
||||
from forms import LoginForm, RegistrationForm
|
||||
from config import Config
|
||||
import logging
|
||||
|
||||
def create_app():
|
||||
"""Application factory pattern"""
|
||||
app = Flask(__name__)
|
||||
app.config.from_object(Config)
|
||||
|
||||
# Initialize extensions
|
||||
db.init_app(app)
|
||||
|
||||
# Setup Flask-Login
|
||||
login_manager = LoginManager()
|
||||
login_manager.init_app(app)
|
||||
login_manager.login_view = 'login'
|
||||
login_manager.login_message = 'Please log in to access this page.'
|
||||
login_manager.login_message_category = 'info'
|
||||
|
||||
@login_manager.user_loader
|
||||
def load_user(user_id):
|
||||
"""Load user for Flask-Login"""
|
||||
return User.query.get(int(user_id))
|
||||
|
||||
# Create database tables
|
||||
with app.app_context():
|
||||
db.create_all()
|
||||
|
||||
return app
|
||||
|
||||
app = create_app()
|
||||
|
||||
@app.route('/')
|
||||
def index():
|
||||
"""Home page"""
|
||||
return render_template('index.html')
|
||||
|
||||
@app.route('/login', methods=['GET', 'POST'])
|
||||
def login():
|
||||
"""User login"""
|
||||
if current_user.is_authenticated:
|
||||
return redirect(url_for('dashboard'))
|
||||
|
||||
form = LoginForm()
|
||||
if form.validate_on_submit():
|
||||
user = User.query.filter_by(username=form.username.data).first()
|
||||
|
||||
if user and user.check_password(form.password.data):
|
||||
login_user(user, remember=form.remember_me.data)
|
||||
next_page = request.args.get('next')
|
||||
flash(f'Welcome back, {user.username}!', 'success')
|
||||
return redirect(next_page) if next_page else redirect(url_for('dashboard'))
|
||||
else:
|
||||
flash('Invalid username or password', 'danger')
|
||||
|
||||
return render_template('login.html', form=form)
|
||||
|
||||
@app.route('/register', methods=['GET', 'POST'])
|
||||
def register():
|
||||
"""User registration"""
|
||||
if current_user.is_authenticated:
|
||||
return redirect(url_for('dashboard'))
|
||||
|
||||
form = RegistrationForm()
|
||||
if form.validate_on_submit():
|
||||
try:
|
||||
user = User(username=form.username.data, email=form.email.data)
|
||||
user.set_password(form.password.data)
|
||||
db.session.add(user)
|
||||
db.session.commit()
|
||||
flash('Registration successful! Please log in.', 'success')
|
||||
return redirect(url_for('login'))
|
||||
except Exception as e:
|
||||
db.session.rollback()
|
||||
flash('Registration failed. Please try again.', 'danger')
|
||||
app.logger.error(f'Registration error: {e}')
|
||||
|
||||
return render_template('register.html', form=form)
|
||||
|
||||
@app.route('/dashboard')
|
||||
@login_required
|
||||
def dashboard():
|
||||
"""Protected user dashboard"""
|
||||
return render_template('dashboard.html')
|
||||
|
||||
@app.route('/logout')
|
||||
@login_required
|
||||
def logout():
|
||||
"""User logout"""
|
||||
username = current_user.username
|
||||
logout_user()
|
||||
flash(f'Goodbye, {username}!', 'info')
|
||||
return redirect(url_for('index'))
|
||||
|
||||
@app.errorhandler(404)
|
||||
def not_found_error(error):
|
||||
"""Handle 404 errors"""
|
||||
return render_template('404.html'), 404
|
||||
|
||||
@app.errorhandler(500)
|
||||
def internal_error(error):
|
||||
"""Handle 500 errors"""
|
||||
db.session.rollback()
|
||||
return render_template('500.html'), 500
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Configure logging for production
|
||||
if not app.debug:
|
||||
import logging
|
||||
from logging.handlers import RotatingFileHandler
|
||||
|
||||
file_handler = RotatingFileHandler('logs/app.log', maxBytes=10240, backupCount=10)
|
||||
file_handler.setFormatter(logging.Formatter(
|
||||
'%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'
|
||||
))
|
||||
file_handler.setLevel(logging.INFO)
|
||||
app.logger.addHandler(file_handler)
|
||||
app.logger.setLevel(logging.INFO)
|
||||
|
||||
app.run(debug=True)
|
||||
```
|
||||
|
||||
## templates/base.html
|
||||
```html
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>{% block title %}Flask Auth App{% endblock %}</title>
|
||||
<link href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css" rel="stylesheet">
|
||||
<link rel="stylesheet" href="{{ url_for('static', filename='style.css') }}">
|
||||
</head>
|
||||
<body>
|
||||
<nav class="navbar navbar-expand-lg navbar-dark bg-dark">
|
||||
<div class="container">
|
||||
<a class="navbar-brand" href="{{ url_for('index') }}">Flask Auth</a>
|
||||
<div class="navbar-nav ms-auto">
|
||||
{% if current_user.is_authenticated %}
|
||||
<a class="nav-link" href="{{ url_for('dashboard') }}">Dashboard</a>
|
||||
<a class="nav-link" href="{{ url_for('logout') }}">Logout ({{ current_user.username }})</a>
|
||||
{% else %}
|
||||
<a class="nav-link" href="{{ url_for('login') }}">Login</a>
|
||||
<a class="nav-link" href="{{ url_for('register') }}">Register</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
</nav>
|
||||
|
||||
<main class="container mt-4">
|
||||
{% with messages = get_flashed_messages(with_categories=true) %}
|
||||
{% if messages %}
|
||||
{% for category, message in messages %}
|
||||
<div class="alert alert-{{ 'danger' if category == 'error' else category }} alert-dismissible fade show">
|
||||
{{ message }}
|
||||
<button type="button" class="btn-close" data-bs-dismiss="alert"></button>
|
||||
</div>
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% endwith %}
|
||||
|
||||
{% block content %}{% endblock %}
|
||||
</main>
|
||||
|
||||
<script src="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/js/bootstrap.bundle.min.js"></script>
|
||||
</body>
|
||||
</html>
|
||||
```
|
||||
|
||||
## templates/index.html
|
||||
```html
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block content %}
|
||||
<div class="row justify-content-center">
|
||||
<div class="col-md-8">
|
||||
<div class="card">
|
||||
<div class="card-body">
|
||||
<h1 class="card-title">Welcome to Flask Auth App</h1>
|
||||
<p class="card-text">A secure authentication system built with Flask.</p>
|
||||
|
||||
{% if current_user.is_authenticated %}
|
||||
<p class="text-success">Hello, {{ current_user.username }}! You are logged in.</p>
|
||||
<a href="{{ url_for('dashboard') }}" class="btn btn-primary">Go to Dashboard</a>
|
||||
{% else %}
|
||||
<p>Please log in to access your dashboard.</p>
|
||||
<a href="{{ url_for('login') }}" class="btn btn-primary me-2">Login</a>
|
||||
<a href="{{ url_for('register') }}" class="btn btn-outline-primary">Register</a>
|
||||
{% endif %}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
```
|
||||
|
||||
## templates/login.html
|
||||
```html
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block title %}Login - Flask Auth App{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="row justify-content-center">
|
||||
<div class="col-md-6">
|
||||
<div class="card">
|
||||
<div class="card-header">
|
||||
<h4 class="mb-0">Sign In</h4>
|
||||
</div>
|
||||
<div class="card-body">
|
||||
<form method="POST">
|
||||
{{ form.hidden_tag() }}
|
||||
|
||||
<div class="mb-3">
|
||||
{{ form.username.label(class="form-label") }}
|
||||
{{ form.username(class="form-control" + (" is-invalid" if form.username.errors else "")) }}
|
||||
{% if form.username.errors %}
|
||||
<div class="invalid-feedback">
|
||||
{% for error in form.username.errors %}{{ error }}{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="mb-3">
|
||||
{{ form.password.label(class="form-label") }}
|
||||
{{ form.password(class="form-control" + (" is-invalid" if form.password.errors else "")) }}
|
||||
{% if form.password.errors %}
|
||||
<div class="invalid-feedback">
|
||||
{% for error in form.password.errors %}{{ error }}{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="mb-3 form-check">
|
||||
{{ form.remember_me(class="form-check-input") }}
|
||||
{{ form.remember_me.label(class="form-check-label") }}
|
||||
</div>
|
||||
|
||||
<div class="d-grid">
|
||||
{{ form.submit(class="btn btn-primary") }}
|
||||
</div>
|
||||
</form>
|
||||
|
||||
<hr>
|
||||
<p class="text-center mb-0">
|
||||
Don't have an account? <a href="{{ url_for('register') }}">Register here</a>
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
```
|
||||
|
||||
## templates/register.html
|
||||
```html
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block title %}Register - Flask Auth App{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="row justify-content-center">
|
||||
<div class="col-md-6">
|
||||
<div class="card">
|
||||
<div class="card-header">
|
||||
<h4 class="mb-0">Create Account</h4>
|
||||
</div>
|
||||
<div class="card-body">
|
||||
<form method="POST">
|
||||
{{ form.hidden_tag() }}
|
||||
|
||||
<div class="mb-3">
|
||||
{{ form.username.label(class="form-label") }}
|
||||
{{ form.username(class="form-control" + (" is-invalid" if form.username.errors else "")) }}
|
||||
{% if form.username.errors %}
|
||||
<div class="invalid-feedback">
|
||||
{% for error in form.username.errors %}{{ error }}{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="mb-3">
|
||||
{{ form.email.label(class="form-label") }}
|
||||
{{ form.email(class="form-control" + (" is-invalid" if form.email.errors else "")) }}
|
||||
{% if form.email.errors %}
|
||||
<div class="invalid-feedback">
|
||||
{% for error in form.email.errors %}{{ error }}{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="mb-3">
|
||||
{{ form.password.label(class="form-label") }}
|
||||
{{ form.password(class="form-control" + (" is-invalid" if form.password.errors else "")) }}
|
||||
{% if form.password.errors %}
|
||||
<div class="invalid-feedback">
|
||||
{% for error in form.password.errors %}{{ error }}{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="mb-3">
|
||||
{{ form.password2.label(class="form-label") }}
|
||||
{{ form.password2(class="form-control" + (" is-invalid" if form.password2.errors else "")) }}
|
||||
{% if form.password2.errors %}
|
||||
<div class="invalid-feedback">
|
||||
{% for error in form.password2.errors %}{{ error }}{% endfor %}
|
||||
</div>
|
||||
{% endif %}
|
||||
</div>
|
||||
|
||||
<div class="d-grid">
|
||||
{{ form.submit(class="btn btn-success") }}
|
||||
</div>
|
||||
</form>
|
||||
|
||||
<hr>
|
||||
<p class="text-center mb-0">
|
||||
Already have an account? <a href="{{ url_for('login') }}">Sign in here</a>
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
```
|
||||
|
||||
## templates/dashboard.html
|
||||
```html
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block title %}Dashboard - Flask Auth App{% endblock %}
|
||||
|
||||
{% block content %}
|
||||
<div class="row">
|
||||
<div class="col-12">
|
||||
<div class="card">
|
||||
<div class="card-header">
|
||||
<h4 class="mb-0">User Dashboard</h4>
|
||||
</div>
|
||||
<div class="card-body">
|
||||
<h5>Welcome, {{ current_user.username }}!</h5>
|
||||
<p class="text-muted">This is your protected dashboard area.</p>
|
||||
|
||||
<div class="row mt-4">
|
||||
<div class="col-md-6">
|
||||
<h6>Account Information</h6>
|
||||
<ul class="list-unstyled">
|
||||
<li><strong>Username:</strong> {{ current_user.username }}</li>
|
||||
<li><strong>Email:</strong> {{ current_user.email }}</li>
|
||||
<li><strong>Member since:</strong> {{ current_user.created_at.strftime('%B %d, %Y') }}</li>
|
||||
</ul>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="mt-4">
|
||||
<a href="{{ url_for('logout') }}" class="btn btn-outline-danger">Logout</a>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
{% endblock %}
|
||||
```
|
||||
|
||||
## static/style.css
|
||||
```css
|
||||
/* Custom styles for Flask Auth App */
|
||||
body {
|
||||
background-color: #f8f9fa;
|
||||
}
|
||||
|
||||
.card {
|
||||
box-shadow: 0 0.125rem 0.25rem rgba(0, 0, 0, 0.075);
|
||||
border: 1px solid rgba(0, 0, 0, 0.125);
|
||||
}
|
||||
|
||||
.navbar-brand {
|
||||
font-weight: bold;
|
||||
}
|
||||
|
||||
.alert {
|
||||
border-radius: 0.375rem;
|
||||
}
|
||||
|
||||
.form-control:focus {
|
||||
border-color: #86b7fe;
|
||||
outline: 0;
|
||||
box-shadow: 0 0 0 0.25rem rgba(13, 110, 253, 0.25);
|
||||
}
|
||||
|
||||
.btn {
|
||||
border-radius: 0.375rem;
|
||||
}
|
||||
|
||||
.card-header {
|
||||
background-color: #f8f9fa;
|
||||
border-bottom: 1px solid rgba(0, 0, 0, 0.125);
|
||||
}
|
||||
|
||||
/* Responsive adjustments */
|
||||
@media (max-width: 768px) {
|
||||
.container {
|
||||
padding-left: 15px;
|
||||
padding-right: 15px;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
This implementation provides:
|
||||
|
||||
- **Secure Authentication**: Password hashing with PBKDF2, CSRF protection, secure sessions
|
||||
- **User Management**: Registration, login, logout with proper validation
|
||||
- **Route Protection**: Login required decorators for protected pages
|
||||
- **Error Handling**: Comprehensive error handling and user feedback
|
||||
- **Production Ready**: Proper configuration management, logging, and security headers
|
||||
- **Responsive UI**: Bootstrap-based responsive design
|
||||
- **Database Integration**: SQLAlchemy ORM with user model
|
||||
|
||||
To run the application:
|
||||
1. Install dependencies: `pip install -r requirements.txt`
|
||||
2. Set environment variables for production
|
||||
3. Run: `python app.py`
|
||||
@@ -0,0 +1,361 @@
|
||||
from flask import Flask, request, jsonify
|
||||
from config import Config
|
||||
from database import init_database
|
||||
from models import db, User
|
||||
from auth import (
|
||||
generate_jwt_token,
|
||||
token_required,
|
||||
validate_email,
|
||||
validate_username,
|
||||
validate_password_strength
|
||||
)
|
||||
import traceback
|
||||
|
||||
app = Flask(__name__)
|
||||
app.config.from_object(Config)
|
||||
|
||||
# Initialize database
|
||||
init_database(app)
|
||||
|
||||
@app.errorhandler(404)
|
||||
def not_found(error):
|
||||
return jsonify({
|
||||
'error': 'Not found',
|
||||
'message': 'The requested resource was not found'
|
||||
}), 404
|
||||
|
||||
@app.errorhandler(500)
|
||||
def internal_error(error):
|
||||
db.session.rollback()
|
||||
return jsonify({
|
||||
'error': 'Internal server error',
|
||||
'message': 'An unexpected error occurred'
|
||||
}), 500
|
||||
|
||||
@app.route('/health', methods=['GET'])
|
||||
def health_check():
|
||||
"""Health check endpoint."""
|
||||
return jsonify({
|
||||
'status': 'healthy',
|
||||
'message': 'Application is running'
|
||||
}), 200
|
||||
|
||||
@app.route('/api/register', methods=['POST'])
|
||||
def register():
|
||||
"""User registration endpoint."""
|
||||
try:
|
||||
data = request.get_json()
|
||||
|
||||
if not data:
|
||||
return jsonify({
|
||||
'error': 'Invalid input',
|
||||
'message': 'JSON data is required'
|
||||
}), 400
|
||||
|
||||
username = data.get('username', '').strip()
|
||||
email = data.get('email', '').strip().lower()
|
||||
password = data.get('password', '')
|
||||
|
||||
# Validate required fields
|
||||
if not username or not email or not password:
|
||||
return jsonify({
|
||||
'error': 'Missing required fields',
|
||||
'message': 'Username, email, and password are required'
|
||||
}), 400
|
||||
|
||||
# Validate username format
|
||||
if not validate_username(username):
|
||||
return jsonify({
|
||||
'error': 'Invalid username',
|
||||
'message': 'Username must be 3-80 characters and contain only letters, numbers, and underscores'
|
||||
}), 400
|
||||
|
||||
# Validate email format
|
||||
if not validate_email(email):
|
||||
return jsonify({
|
||||
'error': 'Invalid email',
|
||||
'message': 'Please provide a valid email address'
|
||||
}), 400
|
||||
|
||||
# Validate password strength
|
||||
is_valid, password_message = validate_password_strength(password)
|
||||
if not is_valid:
|
||||
return jsonify({
|
||||
'error': 'Weak password',
|
||||
'message': password_message
|
||||
}), 400
|
||||
|
||||
# Check if user already exists
|
||||
existing_user = User.query.filter(
|
||||
(User.username == username) | (User.email == email)
|
||||
).first()
|
||||
|
||||
if existing_user:
|
||||
if existing_user.username == username:
|
||||
return jsonify({
|
||||
'error': 'Username taken',
|
||||
'message': 'Username already exists'
|
||||
}), 409
|
||||
else:
|
||||
return jsonify({
|
||||
'error': 'Email taken',
|
||||
'message': 'Email already registered'
|
||||
}), 409
|
||||
|
||||
# Create new user
|
||||
new_user = User(
|
||||
username=username,
|
||||
email=email,
|
||||
password=password
|
||||
)
|
||||
|
||||
db.session.add(new_user)
|
||||
db.session.commit()
|
||||
|
||||
# Generate JWT token
|
||||
token = generate_jwt_token(new_user.id)
|
||||
|
||||
return jsonify({
|
||||
'message': 'User registered successfully',
|
||||
'user': new_user.to_dict(),
|
||||
'token': token
|
||||
}), 201
|
||||
|
||||
except Exception as e:
|
||||
db.session.rollback()
|
||||
app.logger.error(f"Registration error: {str(e)}")
|
||||
app.logger.error(traceback.format_exc())
|
||||
return jsonify({
|
||||
'error': 'Registration failed',
|
||||
'message': 'An error occurred during registration'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/login', methods=['POST'])
|
||||
def login():
|
||||
"""User login endpoint."""
|
||||
try:
|
||||
data = request.get_json()
|
||||
|
||||
if not data:
|
||||
return jsonify({
|
||||
'error': 'Invalid input',
|
||||
'message': 'JSON data is required'
|
||||
}), 400
|
||||
|
||||
login_identifier = data.get('username') or data.get('email', '')
|
||||
password = data.get('password', '')
|
||||
|
||||
if not login_identifier or not password:
|
||||
return jsonify({
|
||||
'error': 'Missing credentials',
|
||||
'message': 'Username/email and password are required'
|
||||
}), 400
|
||||
|
||||
login_identifier = login_identifier.strip().lower()
|
||||
|
||||
# Find user by username or email
|
||||
user = User.query.filter(
|
||||
(User.username == login_identifier) | (User.email == login_identifier)
|
||||
).first()
|
||||
|
||||
if not user or not user.check_password(password):
|
||||
return jsonify({
|
||||
'error': 'Invalid credentials',
|
||||
'message': 'Username/email or password is incorrect'
|
||||
}), 401
|
||||
|
||||
if not user.is_active:
|
||||
return jsonify({
|
||||
'error': 'Account disabled',
|
||||
'message': 'Your account has been disabled'
|
||||
}), 401
|
||||
|
||||
# Generate JWT token
|
||||
token = generate_jwt_token(user.id)
|
||||
|
||||
return jsonify({
|
||||
'message': 'Login successful',
|
||||
'user': user.to_dict(),
|
||||
'token': token
|
||||
}), 200
|
||||
|
||||
except Exception as e:
|
||||
app.logger.error(f"Login error: {str(e)}")
|
||||
app.logger.error(traceback.format_exc())
|
||||
return jsonify({
|
||||
'error': 'Login failed',
|
||||
'message': 'An error occurred during login'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/profile', methods=['GET'])
|
||||
@token_required
|
||||
def get_profile(current_user):
|
||||
"""Get current user profile (protected endpoint)."""
|
||||
return jsonify({
|
||||
'message': 'Profile retrieved successfully',
|
||||
'user': current_user.to_dict()
|
||||
}), 200
|
||||
|
||||
@app.route('/api/profile', methods=['PUT'])
|
||||
@token_required
|
||||
def update_profile(current_user):
|
||||
"""Update current user profile (protected endpoint)."""
|
||||
try:
|
||||
data = request.get_json()
|
||||
|
||||
if not data:
|
||||
return jsonify({
|
||||
'error': 'Invalid input',
|
||||
'message': 'JSON data is required'
|
||||
}), 400
|
||||
|
||||
username = data.get('username')
|
||||
email = data.get('email')
|
||||
|
||||
# Update username if provided
|
||||
if username is not None:
|
||||
username = username.strip()
|
||||
if not validate_username(username):
|
||||
return jsonify({
|
||||
'error': 'Invalid username',
|
||||
'message': 'Username must be 3-80 characters and contain only letters, numbers, and underscores'
|
||||
}), 400
|
||||
|
||||
# Check if username is already taken by another user
|
||||
existing_user = User.query.filter(
|
||||
User.username == username,
|
||||
User.id != current_user.id
|
||||
).first()
|
||||
|
||||
if existing_user:
|
||||
return jsonify({
|
||||
'error': 'Username taken',
|
||||
'message': 'Username already exists'
|
||||
}), 409
|
||||
|
||||
current_user.username = username
|
||||
|
||||
# Update email if provided
|
||||
if email is not None:
|
||||
email = email.strip().lower()
|
||||
if not validate_email(email):
|
||||
return jsonify({
|
||||
'error': 'Invalid email',
|
||||
'message': 'Please provide a valid email address'
|
||||
}), 400
|
||||
|
||||
# Check if email is already taken by another user
|
||||
existing_user = User.query.filter(
|
||||
User.email == email,
|
||||
User.id != current_user.id
|
||||
).first()
|
||||
|
||||
if existing_user:
|
||||
return jsonify({
|
||||
'error': 'Email taken',
|
||||
'message': 'Email already registered'
|
||||
}), 409
|
||||
|
||||
current_user.email = email
|
||||
|
||||
db.session.commit()
|
||||
|
||||
return jsonify({
|
||||
'message': 'Profile updated successfully',
|
||||
'user': current_user.to_dict()
|
||||
}), 200
|
||||
|
||||
except Exception as e:
|
||||
db.session.rollback()
|
||||
app.logger.error(f"Profile update error: {str(e)}")
|
||||
app.logger.error(traceback.format_exc())
|
||||
return jsonify({
|
||||
'error': 'Update failed',
|
||||
'message': 'An error occurred while updating profile'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/change-password', methods=['POST'])
|
||||
@token_required
|
||||
def change_password(current_user):
|
||||
"""Change user password (protected endpoint)."""
|
||||
try:
|
||||
data = request.get_json()
|
||||
|
||||
if not data:
|
||||
return jsonify({
|
||||
'error': 'Invalid input',
|
||||
'message': 'JSON data is required'
|
||||
}), 400
|
||||
|
||||
current_password = data.get('current_password', '')
|
||||
new_password = data.get('new_password', '')
|
||||
|
||||
if not current_password or not new_password:
|
||||
return jsonify({
|
||||
'error': 'Missing password fields',
|
||||
'message': 'Current password and new password are required'
|
||||
}), 400
|
||||
|
||||
# Verify current password
|
||||
if not current_user.check_password(current_password):
|
||||
return jsonify({
|
||||
'error': 'Invalid current password',
|
||||
'message': 'Current password is incorrect'
|
||||
}), 401
|
||||
|
||||
# Validate new password strength
|
||||
is_valid, password_message = validate_password_strength(new_password)
|
||||
if not is_valid:
|
||||
return jsonify({
|
||||
'error': 'Weak password',
|
||||
'message': password_message
|
||||
}), 400
|
||||
|
||||
# Check if new password is different from current
|
||||
if current_user.check_password(new_password):
|
||||
return jsonify({
|
||||
'error': 'Same password',
|
||||
'message': 'New password must be different from current password'
|
||||
}), 400
|
||||
|
||||
# Update password
|
||||
current_user.set_password(new_password)
|
||||
db.session.commit()
|
||||
|
||||
return jsonify({
|
||||
'message': 'Password changed successfully'
|
||||
}), 200
|
||||
|
||||
except Exception as e:
|
||||
db.session.rollback()
|
||||
app.logger.error(f"Password change error: {str(e)}")
|
||||
app.logger.error(traceback.format_exc())
|
||||
return jsonify({
|
||||
'error': 'Password change failed',
|
||||
'message': 'An error occurred while changing password'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/users', methods=['GET'])
|
||||
@token_required
|
||||
def get_users(current_user):
|
||||
"""Get all users (protected endpoint)."""
|
||||
try:
|
||||
users = User.query.filter_by(is_active=True).all()
|
||||
users_data = [user.to_dict() for user in users]
|
||||
|
||||
return jsonify({
|
||||
'message': 'Users retrieved successfully',
|
||||
'users': users_data,
|
||||
'count': len(users_data)
|
||||
}), 200
|
||||
|
||||
except Exception as e:
|
||||
app.logger.error(f"Get users error: {str(e)}")
|
||||
app.logger.error(traceback.format_exc())
|
||||
return jsonify({
|
||||
'error': 'Failed to retrieve users',
|
||||
'message': 'An error occurred while fetching users'
|
||||
}), 500
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(debug=True, host='0.0.0.0', port=5000)
|
||||
@@ -0,0 +1,119 @@
|
||||
import jwt
|
||||
from datetime import datetime, timedelta
|
||||
from functools import wraps
|
||||
from flask import request, jsonify, current_app
|
||||
from models import User
|
||||
import re
|
||||
|
||||
def generate_jwt_token(user_id):
|
||||
"""Generate JWT token for user."""
|
||||
payload = {
|
||||
'user_id': user_id,
|
||||
'exp': datetime.utcnow() + timedelta(hours=current_app.config['JWT_EXPIRATION_HOURS']),
|
||||
'iat': datetime.utcnow()
|
||||
}
|
||||
|
||||
token = jwt.encode(
|
||||
payload,
|
||||
current_app.config['JWT_SECRET_KEY'],
|
||||
algorithm='HS256'
|
||||
)
|
||||
|
||||
return token
|
||||
|
||||
def verify_jwt_token(token):
|
||||
"""Verify JWT token and return user_id if valid."""
|
||||
try:
|
||||
payload = jwt.decode(
|
||||
token,
|
||||
current_app.config['JWT_SECRET_KEY'],
|
||||
algorithms=['HS256']
|
||||
)
|
||||
return payload['user_id']
|
||||
except jwt.ExpiredSignatureError:
|
||||
return None
|
||||
except jwt.InvalidTokenError:
|
||||
return None
|
||||
|
||||
def token_required(f):
|
||||
"""Decorator to require valid JWT token for protected routes."""
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
token = None
|
||||
auth_header = request.headers.get('Authorization')
|
||||
|
||||
if auth_header and auth_header.startswith('Bearer '):
|
||||
token = auth_header.split(' ')[1]
|
||||
|
||||
if not token:
|
||||
return jsonify({
|
||||
'error': 'Token is missing',
|
||||
'message': 'Authorization token is required'
|
||||
}), 401
|
||||
|
||||
user_id = verify_jwt_token(token)
|
||||
if user_id is None:
|
||||
return jsonify({
|
||||
'error': 'Token is invalid or expired',
|
||||
'message': 'Please log in again'
|
||||
}), 401
|
||||
|
||||
current_user = User.query.get(user_id)
|
||||
if not current_user or not current_user.is_active:
|
||||
return jsonify({
|
||||
'error': 'User not found or inactive',
|
||||
'message': 'Invalid user credentials'
|
||||
}), 401
|
||||
|
||||
return f(current_user, *args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
def validate_email(email):
|
||||
"""Validate email format."""
|
||||
pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
|
||||
return re.match(pattern, email) is not None
|
||||
|
||||
def validate_username(username):
|
||||
"""Validate username format."""
|
||||
if not username or len(username) < 3 or len(username) > 80:
|
||||
return False
|
||||
|
||||
pattern = r'^[a-zA-Z0-9_]+$'
|
||||
return re.match(pattern, username) is not None
|
||||
|
||||
def validate_password_strength(password):
|
||||
"""Validate password strength."""
|
||||
min_length = current_app.config.get('PASSWORD_MIN_LENGTH', 8)
|
||||
|
||||
if not password or len(password) < min_length:
|
||||
return False, f'Password must be at least {min_length} characters long'
|
||||
|
||||
if not re.search(r'[A-Z]', password):
|
||||
return False, 'Password must contain at least one uppercase letter'
|
||||
|
||||
if not re.search(r'[a-z]', password):
|
||||
return False, 'Password must contain at least one lowercase letter'
|
||||
|
||||
if not re.search(r'\d', password):
|
||||
return False, 'Password must contain at least one digit'
|
||||
|
||||
if not re.search(r'[!@#$%^&*(),.?":{}|<>]', password):
|
||||
return False, 'Password must contain at least one special character'
|
||||
|
||||
return True, 'Password is valid'
|
||||
|
||||
def get_current_user_from_token():
|
||||
"""Extract current user from request token."""
|
||||
auth_header = request.headers.get('Authorization')
|
||||
|
||||
if not auth_header or not auth_header.startswith('Bearer '):
|
||||
return None
|
||||
|
||||
token = auth_header.split(' ')[1]
|
||||
user_id = verify_jwt_token(token)
|
||||
|
||||
if user_id:
|
||||
return User.query.get(user_id)
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,12 @@
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
class Config:
|
||||
SECRET_KEY = os.environ.get('SECRET_KEY') or 'dev-secret-key-change-in-production'
|
||||
SQLALCHEMY_DATABASE_URI = os.environ.get('DATABASE_URL') or 'sqlite:///app.db'
|
||||
SQLALCHEMY_TRACK_MODIFICATIONS = False
|
||||
JWT_SECRET_KEY = os.environ.get('JWT_SECRET_KEY') or 'jwt-secret-key-change-in-production'
|
||||
JWT_EXPIRATION_HOURS = int(os.environ.get('JWT_EXPIRATION_HOURS', 24))
|
||||
PASSWORD_MIN_LENGTH = int(os.environ.get('PASSWORD_MIN_LENGTH', 8))
|
||||
@@ -0,0 +1,36 @@
|
||||
from flask import Flask
|
||||
from models import db, User
|
||||
import os
|
||||
|
||||
def create_app_context(app):
|
||||
"""Create application context and initialize database."""
|
||||
with app.app_context():
|
||||
db.create_all()
|
||||
return app
|
||||
|
||||
def init_database(app):
|
||||
"""Initialize database with Flask app."""
|
||||
db.init_app(app)
|
||||
|
||||
# Create tables if they don't exist
|
||||
with app.app_context():
|
||||
db.create_all()
|
||||
|
||||
# Create default admin user if it doesn't exist
|
||||
admin_user = User.query.filter_by(username='admin').first()
|
||||
if not admin_user and os.environ.get('CREATE_ADMIN_USER', 'false').lower() == 'true':
|
||||
admin_email = os.environ.get('ADMIN_EMAIL', 'admin@example.com')
|
||||
admin_password = os.environ.get('ADMIN_PASSWORD', 'admin123')
|
||||
|
||||
admin = User(
|
||||
username='admin',
|
||||
email=admin_email,
|
||||
password=admin_password
|
||||
)
|
||||
db.session.add(admin)
|
||||
db.session.commit()
|
||||
print("Admin user created successfully")
|
||||
|
||||
def get_db():
|
||||
"""Get database instance."""
|
||||
return db
|
||||
@@ -0,0 +1,45 @@
|
||||
from flask_sqlalchemy import SQLAlchemy
|
||||
from datetime import datetime
|
||||
import bcrypt
|
||||
|
||||
db = SQLAlchemy()
|
||||
|
||||
class User(db.Model):
|
||||
__tablename__ = 'users'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
username = db.Column(db.String(80), unique=True, nullable=False, index=True)
|
||||
email = db.Column(db.String(120), unique=True, nullable=False, index=True)
|
||||
password_hash = db.Column(db.String(128), nullable=False)
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
is_active = db.Column(db.Boolean, default=True)
|
||||
|
||||
def __init__(self, username, email, password):
|
||||
self.username = username
|
||||
self.email = email
|
||||
self.set_password(password)
|
||||
|
||||
def set_password(self, password):
|
||||
"""Hash and set the user's password."""
|
||||
password_bytes = password.encode('utf-8')
|
||||
salt = bcrypt.gensalt()
|
||||
self.password_hash = bcrypt.hashpw(password_bytes, salt).decode('utf-8')
|
||||
|
||||
def check_password(self, password):
|
||||
"""Check if the provided password matches the stored hash."""
|
||||
password_bytes = password.encode('utf-8')
|
||||
hash_bytes = self.password_hash.encode('utf-8')
|
||||
return bcrypt.checkpw(password_bytes, hash_bytes)
|
||||
|
||||
def to_dict(self):
|
||||
"""Convert user object to dictionary (excluding sensitive data)."""
|
||||
return {
|
||||
'id': self.id,
|
||||
'username': self.username,
|
||||
'email': self.email,
|
||||
'created_at': self.created_at.isoformat(),
|
||||
'is_active': self.is_active
|
||||
}
|
||||
|
||||
def __repr__(self):
|
||||
return f'<User {self.username}>'
|
||||
@@ -0,0 +1,5 @@
|
||||
flask>=2.3.0
|
||||
flask-sqlalchemy>=3.0.0
|
||||
bcrypt>=4.0.0
|
||||
pyjwt>=2.8.0
|
||||
python-dotenv>=1.0.0
|
||||
@@ -0,0 +1 @@
|
||||
requests>=2.28.0
|
||||
@@ -0,0 +1,193 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Weather CLI application using OpenWeatherMap API
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
@dataclass
|
||||
class WeatherData:
|
||||
"""Data class to hold weather information"""
|
||||
temperature: float
|
||||
humidity: int
|
||||
wind_speed: float
|
||||
description: str
|
||||
city: str
|
||||
country: str
|
||||
|
||||
|
||||
class WeatherAPIClient:
|
||||
"""Client for OpenWeatherMap API"""
|
||||
|
||||
BASE_URL = "https://api.openweathermap.org/data/2.5/weather"
|
||||
TIMEOUT = 10
|
||||
|
||||
def __init__(self, api_key: str):
|
||||
self.api_key = api_key
|
||||
self.session = requests.Session()
|
||||
|
||||
def get_weather(self, city: str) -> WeatherData:
|
||||
"""Fetch weather data for a given city"""
|
||||
params = {
|
||||
'q': city,
|
||||
'appid': self.api_key,
|
||||
'units': 'metric'
|
||||
}
|
||||
|
||||
try:
|
||||
response = self.session.get(
|
||||
self.BASE_URL,
|
||||
params=params,
|
||||
timeout=self.TIMEOUT
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
raise WeatherAPIError("Request timed out. Please try again.")
|
||||
except requests.exceptions.ConnectionError:
|
||||
raise WeatherAPIError("Unable to connect to weather service. Check your internet connection.")
|
||||
except requests.exceptions.HTTPError as e:
|
||||
if response.status_code == 401:
|
||||
raise WeatherAPIError("Invalid API key. Please check your OpenWeatherMap API key.")
|
||||
elif response.status_code == 404:
|
||||
raise WeatherAPIError(f"City '{city}' not found. Please check the spelling.")
|
||||
elif response.status_code == 429:
|
||||
raise WeatherAPIError("API rate limit exceeded. Please try again later.")
|
||||
else:
|
||||
raise WeatherAPIError(f"API request failed with status {response.status_code}")
|
||||
except requests.exceptions.RequestException:
|
||||
raise WeatherAPIError("An unexpected error occurred while fetching weather data.")
|
||||
|
||||
try:
|
||||
data = response.json()
|
||||
return self._parse_weather_data(data)
|
||||
except json.JSONDecodeError:
|
||||
raise WeatherAPIError("Invalid response format from weather service.")
|
||||
except KeyError as e:
|
||||
raise WeatherAPIError(f"Unexpected response format: missing {e}")
|
||||
|
||||
def _parse_weather_data(self, data: dict) -> WeatherData:
|
||||
"""Parse API response into WeatherData object"""
|
||||
try:
|
||||
return WeatherData(
|
||||
temperature=data['main']['temp'],
|
||||
humidity=data['main']['humidity'],
|
||||
wind_speed=data.get('wind', {}).get('speed', 0.0),
|
||||
description=data['weather'][0]['description'].title(),
|
||||
city=data['name'],
|
||||
country=data['sys']['country']
|
||||
)
|
||||
except (KeyError, IndexError, TypeError) as e:
|
||||
raise WeatherAPIError(f"Failed to parse weather data: {e}")
|
||||
|
||||
|
||||
class WeatherAPIError(Exception):
|
||||
"""Custom exception for weather API errors"""
|
||||
pass
|
||||
|
||||
|
||||
class WeatherFormatter:
|
||||
"""Format weather data for terminal output"""
|
||||
|
||||
@staticmethod
|
||||
def format_weather(weather: WeatherData) -> str:
|
||||
"""Format weather data for display"""
|
||||
return f"""
|
||||
Weather for {weather.city}, {weather.country}
|
||||
{'=' * (len(weather.city) + len(weather.country) + 12)}
|
||||
|
||||
Temperature: {weather.temperature:.1f}°C
|
||||
Humidity: {weather.humidity}%
|
||||
Wind Speed: {weather.wind_speed:.1f} m/s
|
||||
Conditions: {weather.description}
|
||||
"""
|
||||
|
||||
|
||||
def get_api_key() -> Optional[str]:
|
||||
"""Get API key from environment variable"""
|
||||
return os.environ.get('OPENWEATHERMAP_API_KEY')
|
||||
|
||||
|
||||
def create_parser() -> argparse.ArgumentParser:
|
||||
"""Create and configure argument parser"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Get current weather information for any city',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
%(prog)s "New York"
|
||||
%(prog)s London
|
||||
%(prog)s "San Francisco" --api-key YOUR_API_KEY
|
||||
|
||||
Environment Variables:
|
||||
OPENWEATHERMAP_API_KEY Your OpenWeatherMap API key
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'city',
|
||||
help='City name to get weather for (use quotes for multi-word cities)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--api-key',
|
||||
help='OpenWeatherMap API key (overrides environment variable)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--version',
|
||||
action='version',
|
||||
version='Weather CLI 1.0.0'
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Main application entry point"""
|
||||
parser = create_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key from command line or environment
|
||||
api_key = args.api_key or get_api_key()
|
||||
|
||||
if not api_key:
|
||||
print("Error: OpenWeatherMap API key required.", file=sys.stderr)
|
||||
print("Either set OPENWEATHERMAP_API_KEY environment variable", file=sys.stderr)
|
||||
print("or provide --api-key argument.", file=sys.stderr)
|
||||
print("\nGet your free API key at: https://openweathermap.org/api", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
try:
|
||||
# Initialize API client and fetch weather
|
||||
client = WeatherAPIClient(api_key)
|
||||
weather_data = client.get_weather(args.city)
|
||||
|
||||
# Format and display results
|
||||
formatter = WeatherFormatter()
|
||||
output = formatter.format_weather(weather_data)
|
||||
print(output)
|
||||
|
||||
return 0
|
||||
|
||||
except WeatherAPIError as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
return 1
|
||||
except KeyboardInterrupt:
|
||||
print("\nOperation cancelled by user.", file=sys.stderr)
|
||||
return 1
|
||||
except Exception as e:
|
||||
print(f"Unexpected error: {e}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,121 @@
|
||||
from flask import Flask, request, jsonify, render_template, redirect
|
||||
import validators
|
||||
import string
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# In-memory storage
|
||||
url_mapping = {} # counter -> original_url
|
||||
counter = 1
|
||||
|
||||
# Base62 characters for encoding
|
||||
BASE62_CHARS = string.ascii_lowercase + string.ascii_uppercase + string.digits
|
||||
|
||||
def encode_base62(num):
|
||||
"""Convert integer to base62 string"""
|
||||
if num == 0:
|
||||
return BASE62_CHARS[0]
|
||||
|
||||
result = []
|
||||
while num > 0:
|
||||
result.append(BASE62_CHARS[num % 62])
|
||||
num //= 62
|
||||
|
||||
return ''.join(reversed(result))
|
||||
|
||||
def decode_base62(encoded):
|
||||
"""Convert base62 string back to integer"""
|
||||
result = 0
|
||||
for char in encoded:
|
||||
if char in BASE62_CHARS:
|
||||
result = result * 62 + BASE62_CHARS.index(char)
|
||||
else:
|
||||
return None
|
||||
return result
|
||||
|
||||
@app.route('/')
|
||||
def index():
|
||||
"""Serve the landing page"""
|
||||
return render_template('index.html')
|
||||
|
||||
@app.route('/shorten', methods=['POST'])
|
||||
def shorten_url():
|
||||
"""Shorten a URL and return the short code"""
|
||||
global counter
|
||||
|
||||
# Get URL from JSON request or form data
|
||||
if request.is_json:
|
||||
data = request.get_json()
|
||||
if not data or 'url' not in data:
|
||||
return jsonify({'error': 'URL is required'}), 400
|
||||
url = data['url']
|
||||
else:
|
||||
url = request.form.get('url')
|
||||
if not url:
|
||||
return jsonify({'error': 'URL is required'}), 400
|
||||
|
||||
# Validate URL
|
||||
if not validators.url(url):
|
||||
return jsonify({'error': 'Invalid URL format'}), 400
|
||||
|
||||
# Check if URL already exists
|
||||
for existing_counter, existing_url in url_mapping.items():
|
||||
if existing_url == url:
|
||||
short_code = encode_base62(existing_counter)
|
||||
if request.is_json:
|
||||
return jsonify({
|
||||
'short_code': short_code,
|
||||
'short_url': request.host_url + short_code,
|
||||
'original_url': url
|
||||
})
|
||||
else:
|
||||
return render_template('index.html',
|
||||
short_code=short_code,
|
||||
short_url=request.host_url + short_code,
|
||||
original_url=url)
|
||||
|
||||
# Store new URL
|
||||
url_mapping[counter] = url
|
||||
short_code = encode_base62(counter)
|
||||
counter += 1
|
||||
|
||||
if request.is_json:
|
||||
return jsonify({
|
||||
'short_code': short_code,
|
||||
'short_url': request.host_url + short_code,
|
||||
'original_url': url
|
||||
})
|
||||
else:
|
||||
return render_template('index.html',
|
||||
short_code=short_code,
|
||||
short_url=request.host_url + short_code,
|
||||
original_url=url)
|
||||
|
||||
@app.route('/<short_code>')
|
||||
def redirect_url(short_code):
|
||||
"""Redirect to original URL using short code"""
|
||||
# Decode the short code
|
||||
counter_id = decode_base62(short_code)
|
||||
|
||||
if counter_id is None:
|
||||
return jsonify({'error': 'Invalid short code format'}), 404
|
||||
|
||||
# Look up the original URL
|
||||
if counter_id not in url_mapping:
|
||||
return jsonify({'error': 'Short code not found'}), 404
|
||||
|
||||
original_url = url_mapping[counter_id]
|
||||
return redirect(original_url, code=302)
|
||||
|
||||
@app.errorhandler(404)
|
||||
def not_found(error):
|
||||
"""Handle 404 errors"""
|
||||
return jsonify({'error': 'Endpoint not found'}), 404
|
||||
|
||||
@app.errorhandler(500)
|
||||
def internal_error(error):
|
||||
"""Handle 500 errors"""
|
||||
return jsonify({'error': 'Internal server error'}), 500
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(debug=True, host='0.0.0.0', port=5000)
|
||||
@@ -0,0 +1,2 @@
|
||||
flask>=2.0.0
|
||||
validators>=0.20.0
|
||||
@@ -0,0 +1,161 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>URL Shortener</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: Arial, sans-serif;
|
||||
max-width: 800px;
|
||||
margin: 0 auto;
|
||||
padding: 20px;
|
||||
background-color: #f5f5f5;
|
||||
}
|
||||
.container {
|
||||
background-color: white;
|
||||
padding: 30px;
|
||||
border-radius: 10px;
|
||||
box-shadow: 0 2px 10px rgba(0,0,0,0.1);
|
||||
}
|
||||
h1 {
|
||||
color: #333;
|
||||
text-align: center;
|
||||
margin-bottom: 30px;
|
||||
}
|
||||
.form-group {
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
label {
|
||||
display: block;
|
||||
margin-bottom: 5px;
|
||||
font-weight: bold;
|
||||
color: #555;
|
||||
}
|
||||
input[type="url"] {
|
||||
width: 100%;
|
||||
padding: 12px;
|
||||
border: 2px solid #ddd;
|
||||
border-radius: 5px;
|
||||
font-size: 16px;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
input[type="url"]:focus {
|
||||
border-color: #007bff;
|
||||
outline: none;
|
||||
}
|
||||
button {
|
||||
background-color: #007bff;
|
||||
color: white;
|
||||
padding: 12px 30px;
|
||||
border: none;
|
||||
border-radius: 5px;
|
||||
font-size: 16px;
|
||||
cursor: pointer;
|
||||
width: 100%;
|
||||
}
|
||||
button:hover {
|
||||
background-color: #0056b3;
|
||||
}
|
||||
.result {
|
||||
margin-top: 30px;
|
||||
padding: 20px;
|
||||
background-color: #d4edda;
|
||||
border: 1px solid #c3e6cb;
|
||||
border-radius: 5px;
|
||||
}
|
||||
.result h2 {
|
||||
color: #155724;
|
||||
margin-top: 0;
|
||||
}
|
||||
.result p {
|
||||
margin: 10px 0;
|
||||
color: #155724;
|
||||
}
|
||||
.short-url {
|
||||
font-weight: bold;
|
||||
word-break: break-all;
|
||||
}
|
||||
.copy-btn {
|
||||
background-color: #28a745;
|
||||
padding: 8px 15px;
|
||||
font-size: 14px;
|
||||
margin-top: 10px;
|
||||
width: auto;
|
||||
}
|
||||
.copy-btn:hover {
|
||||
background-color: #218838;
|
||||
}
|
||||
.footer {
|
||||
text-align: center;
|
||||
margin-top: 30px;
|
||||
color: #666;
|
||||
font-size: 14px;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<h1>🔗 URL Shortener</h1>
|
||||
|
||||
<form method="POST" action="/shorten">
|
||||
<div class="form-group">
|
||||
<label for="url">Enter URL to shorten:</label>
|
||||
<input type="url" id="url" name="url" placeholder="https://example.com" required>
|
||||
</div>
|
||||
<button type="submit">Shorten URL</button>
|
||||
</form>
|
||||
|
||||
{% if short_code %}
|
||||
<div class="result">
|
||||
<h2>✅ Success!</h2>
|
||||
<p><strong>Original URL:</strong> {{ original_url }}</p>
|
||||
<p><strong>Short URL:</strong> <span class="short-url" id="shortUrl">{{ short_url }}</span></p>
|
||||
<p><strong>Short Code:</strong> {{ short_code }}</p>
|
||||
<button class="copy-btn" onclick="copyToClipboard()">📋 Copy Short URL</button>
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<div class="footer">
|
||||
<p>Enter any valid URL above to generate a shortened version.</p>
|
||||
<p>The shortened URL will redirect to your original link.</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
function copyToClipboard() {
|
||||
const shortUrl = document.getElementById('shortUrl').textContent;
|
||||
navigator.clipboard.writeText(shortUrl).then(function() {
|
||||
const btn = document.querySelector('.copy-btn');
|
||||
const originalText = btn.textContent;
|
||||
btn.textContent = '✅ Copied!';
|
||||
btn.style.backgroundColor = '#28a745';
|
||||
setTimeout(function() {
|
||||
btn.textContent = originalText;
|
||||
btn.style.backgroundColor = '#28a745';
|
||||
}, 2000);
|
||||
}).catch(function(err) {
|
||||
console.error('Failed to copy: ', err);
|
||||
// Fallback for older browsers
|
||||
const textArea = document.createElement('textarea');
|
||||
textArea.value = shortUrl;
|
||||
document.body.appendChild(textArea);
|
||||
textArea.focus();
|
||||
textArea.select();
|
||||
try {
|
||||
document.execCommand('copy');
|
||||
const btn = document.querySelector('.copy-btn');
|
||||
const originalText = btn.textContent;
|
||||
btn.textContent = '✅ Copied!';
|
||||
setTimeout(function() {
|
||||
btn.textContent = originalText;
|
||||
}, 2000);
|
||||
} catch (err) {
|
||||
console.error('Fallback copy failed: ', err);
|
||||
}
|
||||
document.body.removeChild(textArea);
|
||||
});
|
||||
}
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,376 @@
|
||||
import os
|
||||
import argparse
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Optional
|
||||
import numpy as np
|
||||
import faiss
|
||||
import tiktoken
|
||||
import openai
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RAGApplication:
|
||||
"""
|
||||
A complete RAG (Retrieval-Augmented Generation) application that processes documents,
|
||||
creates embeddings, stores them in FAISS, and answers questions using OpenAI GPT.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "gpt-3.5-turbo", embedding_model: str = "text-embedding-ada-002"):
|
||||
"""
|
||||
Initialize the RAG application.
|
||||
|
||||
Args:
|
||||
model_name: OpenAI model name for chat completions
|
||||
embedding_model: OpenAI model name for embeddings
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.embedding_model = embedding_model
|
||||
self.chunk_size = 1000
|
||||
self.chunk_overlap = 200
|
||||
self.max_chunks_per_query = 5
|
||||
|
||||
# Initialize OpenAI client
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("OPENAI_API_KEY environment variable not found")
|
||||
|
||||
self.client = openai.OpenAI(api_key=api_key)
|
||||
|
||||
# Initialize tokenizer
|
||||
try:
|
||||
self.tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
|
||||
except KeyError:
|
||||
self.tokenizer = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
# Storage for documents and embeddings
|
||||
self.documents: List[str] = []
|
||||
self.chunks: List[str] = []
|
||||
self.embeddings: Optional[np.ndarray] = None
|
||||
self.faiss_index: Optional[faiss.Index] = None
|
||||
|
||||
def load_documents(self, docs_dir: str) -> None:
|
||||
"""
|
||||
Load all .txt files from the specified directory.
|
||||
|
||||
Args:
|
||||
docs_dir: Path to directory containing .txt files
|
||||
"""
|
||||
docs_path = Path(docs_dir)
|
||||
|
||||
if not docs_path.exists():
|
||||
raise FileNotFoundError(f"Directory not found: {docs_dir}")
|
||||
|
||||
if not docs_path.is_dir():
|
||||
raise ValueError(f"Path is not a directory: {docs_dir}")
|
||||
|
||||
txt_files = list(docs_path.glob("*.txt"))
|
||||
|
||||
if not txt_files:
|
||||
raise ValueError(f"No .txt files found in directory: {docs_dir}")
|
||||
|
||||
logger.info(f"Found {len(txt_files)} .txt files in {docs_dir}")
|
||||
|
||||
for txt_file in txt_files:
|
||||
try:
|
||||
with open(txt_file, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
if content.strip(): # Only add non-empty files
|
||||
self.documents.append(content)
|
||||
logger.info(f"Loaded: {txt_file.name}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading {txt_file}: {e}")
|
||||
continue
|
||||
|
||||
if not self.documents:
|
||||
raise ValueError("No valid documents were loaded")
|
||||
|
||||
logger.info(f"Successfully loaded {len(self.documents)} documents")
|
||||
|
||||
def chunk_text(self, text: str) -> List[str]:
|
||||
"""
|
||||
Split text into chunks with overlap for better context preservation.
|
||||
|
||||
Args:
|
||||
text: Text to chunk
|
||||
|
||||
Returns:
|
||||
List of text chunks
|
||||
"""
|
||||
tokens = self.tokenizer.encode(text)
|
||||
chunks = []
|
||||
|
||||
for i in range(0, len(tokens), self.chunk_size - self.chunk_overlap):
|
||||
chunk_tokens = tokens[i:i + self.chunk_size]
|
||||
chunk_text = self.tokenizer.decode(chunk_tokens)
|
||||
chunks.append(chunk_text)
|
||||
|
||||
return chunks
|
||||
|
||||
def process_documents(self) -> None:
|
||||
"""
|
||||
Process all loaded documents by chunking them into smaller pieces.
|
||||
"""
|
||||
if not self.documents:
|
||||
raise ValueError("No documents loaded. Call load_documents() first.")
|
||||
|
||||
logger.info("Processing documents into chunks...")
|
||||
|
||||
self.chunks = []
|
||||
for i, doc in enumerate(self.documents):
|
||||
doc_chunks = self.chunk_text(doc)
|
||||
self.chunks.extend(doc_chunks)
|
||||
logger.info(f"Document {i+1}: created {len(doc_chunks)} chunks")
|
||||
|
||||
logger.info(f"Total chunks created: {len(self.chunks)}")
|
||||
|
||||
def get_embeddings(self, texts: List[str], batch_size: int = 100) -> np.ndarray:
|
||||
"""
|
||||
Get embeddings for a list of texts using OpenAI API.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed
|
||||
batch_size: Number of texts to process in each API call
|
||||
|
||||
Returns:
|
||||
NumPy array of embeddings
|
||||
"""
|
||||
embeddings = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i:i + batch_size]
|
||||
|
||||
try:
|
||||
response = self.client.embeddings.create(
|
||||
model=self.embedding_model,
|
||||
input=batch
|
||||
)
|
||||
|
||||
batch_embeddings = [data.embedding for data in response.data]
|
||||
embeddings.extend(batch_embeddings)
|
||||
|
||||
logger.info(f"Generated embeddings for batch {i//batch_size + 1}/{(len(texts)-1)//batch_size + 1}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating embeddings for batch {i//batch_size + 1}: {e}")
|
||||
raise
|
||||
|
||||
return np.array(embeddings, dtype=np.float32)
|
||||
|
||||
def create_vector_store(self) -> None:
|
||||
"""
|
||||
Create embeddings for all chunks and initialize FAISS index.
|
||||
"""
|
||||
if not self.chunks:
|
||||
raise ValueError("No chunks available. Call process_documents() first.")
|
||||
|
||||
logger.info("Generating embeddings for chunks...")
|
||||
self.embeddings = self.get_embeddings(self.chunks)
|
||||
|
||||
# Initialize FAISS index
|
||||
embedding_dim = self.embeddings.shape[1]
|
||||
self.faiss_index = faiss.IndexFlatIP(embedding_dim) # Inner product for cosine similarity
|
||||
|
||||
# Normalize embeddings for cosine similarity
|
||||
faiss.normalize_L2(self.embeddings)
|
||||
|
||||
# Add embeddings to index
|
||||
self.faiss_index.add(self.embeddings)
|
||||
|
||||
logger.info(f"Created FAISS index with {self.faiss_index.ntotal} vectors")
|
||||
|
||||
def retrieve_relevant_chunks(self, query: str, k: int = None) -> List[Tuple[str, float]]:
|
||||
"""
|
||||
Retrieve the most relevant chunks for a given query.
|
||||
|
||||
Args:
|
||||
query: User query
|
||||
k: Number of chunks to retrieve (default: self.max_chunks_per_query)
|
||||
|
||||
Returns:
|
||||
List of tuples (chunk_text, similarity_score)
|
||||
"""
|
||||
if self.faiss_index is None:
|
||||
raise ValueError("Vector store not initialized. Call create_vector_store() first.")
|
||||
|
||||
if k is None:
|
||||
k = min(self.max_chunks_per_query, len(self.chunks))
|
||||
|
||||
# Get query embedding
|
||||
query_embedding = self.get_embeddings([query])
|
||||
faiss.normalize_L2(query_embedding)
|
||||
|
||||
# Search for similar chunks
|
||||
similarities, indices = self.faiss_index.search(query_embedding, k)
|
||||
|
||||
# Return chunks with their similarity scores
|
||||
results = []
|
||||
for i, (similarity, idx) in enumerate(zip(similarities[0], indices[0])):
|
||||
if idx < len(self.chunks): # Ensure valid index
|
||||
results.append((self.chunks[idx], float(similarity)))
|
||||
|
||||
return results
|
||||
|
||||
def generate_answer(self, query: str, context_chunks: List[str]) -> str:
|
||||
"""
|
||||
Generate an answer using OpenAI GPT with the retrieved context.
|
||||
|
||||
Args:
|
||||
query: User query
|
||||
context_chunks: List of relevant text chunks
|
||||
|
||||
Returns:
|
||||
Generated answer
|
||||
"""
|
||||
# Prepare context
|
||||
context = "\n\n".join(context_chunks)
|
||||
|
||||
# Create prompt
|
||||
prompt = f"""Based on the following context, please answer the question. If the answer cannot be found in the context, please say so.
|
||||
|
||||
Context:
|
||||
{context}
|
||||
|
||||
Question: {query}
|
||||
|
||||
Answer:"""
|
||||
|
||||
try:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant that answers questions based on provided context."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
max_tokens=1000,
|
||||
temperature=0.1
|
||||
)
|
||||
|
||||
return response.choices[0].message.content.strip()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating answer: {e}")
|
||||
return f"Sorry, I encountered an error while generating the answer: {e}"
|
||||
|
||||
def answer_question(self, question: str) -> str:
|
||||
"""
|
||||
Complete pipeline to answer a question using RAG.
|
||||
|
||||
Args:
|
||||
question: User question
|
||||
|
||||
Returns:
|
||||
Generated answer with context information
|
||||
"""
|
||||
logger.info(f"Processing question: {question}")
|
||||
|
||||
# Retrieve relevant chunks
|
||||
relevant_chunks = self.retrieve_relevant_chunks(question)
|
||||
|
||||
if not relevant_chunks:
|
||||
return "Sorry, I couldn't find any relevant information to answer your question."
|
||||
|
||||
# Extract chunk texts and log retrieval info
|
||||
chunk_texts = [chunk for chunk, score in relevant_chunks]
|
||||
logger.info(f"Retrieved {len(chunk_texts)} relevant chunks")
|
||||
|
||||
# Generate answer
|
||||
answer = self.generate_answer(question, chunk_texts)
|
||||
|
||||
# Add metadata about sources
|
||||
source_info = f"\n\n[Answer based on {len(chunk_texts)} relevant text segments from the document collection]"
|
||||
|
||||
return answer + source_info
|
||||
|
||||
def initialize(self, docs_dir: str) -> None:
|
||||
"""
|
||||
Complete initialization of the RAG system.
|
||||
|
||||
Args:
|
||||
docs_dir: Directory containing .txt files
|
||||
"""
|
||||
logger.info("Initializing RAG system...")
|
||||
|
||||
self.load_documents(docs_dir)
|
||||
self.process_documents()
|
||||
self.create_vector_store()
|
||||
|
||||
logger.info("RAG system initialized successfully!")
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Command-line interface for the RAG application.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="RAG Question Answering System")
|
||||
parser.add_argument(
|
||||
"--docs-dir",
|
||||
required=True,
|
||||
help="Directory containing .txt documents to index"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--question",
|
||||
help="Question to ask (if not provided, will enter interactive mode)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="gpt-3.5-turbo",
|
||||
help="OpenAI model name for answer generation"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
# Initialize RAG application
|
||||
rag = RAGApplication(model_name=args.model)
|
||||
rag.initialize(args.docs_dir)
|
||||
|
||||
if args.question:
|
||||
# Single question mode
|
||||
answer = rag.answer_question(args.question)
|
||||
print("\nQuestion:", args.question)
|
||||
print("\nAnswer:", answer)
|
||||
else:
|
||||
# Interactive mode
|
||||
print("\nRAG system ready! Type 'quit' or 'exit' to stop.")
|
||||
print("=" * 50)
|
||||
|
||||
while True:
|
||||
try:
|
||||
question = input("\nEnter your question: ").strip()
|
||||
|
||||
if question.lower() in ['quit', 'exit', 'q']:
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
if not question:
|
||||
print("Please enter a question.")
|
||||
continue
|
||||
|
||||
answer = rag.answer_question(question)
|
||||
print(f"\nAnswer: {answer}")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nGoodbye!")
|
||||
break
|
||||
except EOFError:
|
||||
print("\n\nGoodbye!")
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Application error: {e}")
|
||||
print(f"Error: {e}")
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
@@ -0,0 +1,5 @@
|
||||
openai>=1.0.0
|
||||
faiss-cpu>=1.7.0
|
||||
tiktoken>=0.5.0
|
||||
numpy>=1.21.0
|
||||
python-dotenv>=1.0.0
|
||||
@@ -0,0 +1,6 @@
|
||||
# OpenAI API Configuration
|
||||
OPENAI_API_KEY=your_openai_api_key_here
|
||||
|
||||
# Optional: Set default model preferences
|
||||
# EMBEDDING_MODEL=text-embedding-3-small
|
||||
# CHAT_MODEL=gpt-3.5-turbo
|
||||
@@ -0,0 +1,73 @@
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass
|
||||
class RAGConfig:
|
||||
"""Configuration for the RAG system."""
|
||||
|
||||
# OpenAI API settings
|
||||
openai_api_key: str = ""
|
||||
embedding_model: str = "text-embedding-3-small"
|
||||
chat_model: str = "gpt-3.5-turbo"
|
||||
temperature: float = 0.1
|
||||
|
||||
# Chunking parameters
|
||||
chunk_size: int = 1000
|
||||
chunk_overlap: int = 200
|
||||
|
||||
# Embedding settings
|
||||
embedding_batch_size: int = 100
|
||||
embedding_dimension: int = 1536
|
||||
|
||||
# Retrieval settings
|
||||
top_k: int = 5
|
||||
similarity_threshold: float = 0.7
|
||||
|
||||
# File paths
|
||||
documents_dir: str = "documents"
|
||||
index_dir: str = "index"
|
||||
|
||||
def get_index_path(self) -> str:
|
||||
"""Get path to FAISS index file."""
|
||||
return str(Path(self.index_dir) / "faiss.index")
|
||||
|
||||
def get_metadata_path(self) -> str:
|
||||
"""Get path to metadata file."""
|
||||
return str(Path(self.index_dir) / "metadata.pkl")
|
||||
|
||||
def ensure_directories(self):
|
||||
"""Create necessary directories if they don't exist."""
|
||||
Path(self.documents_dir).mkdir(exist_ok=True)
|
||||
Path(self.index_dir).mkdir(exist_ok=True)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""Configuration for different model options."""
|
||||
|
||||
# Available embedding models
|
||||
EMBEDDING_MODELS = {
|
||||
"text-embedding-3-small": 1536,
|
||||
"text-embedding-3-large": 3072,
|
||||
"text-embedding-ada-002": 1536
|
||||
}
|
||||
|
||||
# Available chat models
|
||||
CHAT_MODELS = [
|
||||
"gpt-3.5-turbo",
|
||||
"gpt-4",
|
||||
"gpt-4-turbo-preview"
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def get_embedding_dimension(cls, model: str) -> int:
|
||||
"""Get embedding dimension for a model."""
|
||||
return cls.EMBEDDING_MODELS.get(model, 1536)
|
||||
|
||||
@classmethod
|
||||
def validate_models(cls, embedding_model: str, chat_model: str) -> bool:
|
||||
"""Validate model selections."""
|
||||
embedding_valid = embedding_model in cls.EMBEDDING_MODELS
|
||||
chat_valid = chat_model in cls.CHAT_MODELS
|
||||
return embedding_valid and chat_valid
|
||||
@@ -0,0 +1,228 @@
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from config import RAGConfig, ModelConfig
|
||||
from src.document_loader import DocumentLoader
|
||||
from src.chunker import TextChunker
|
||||
from src.embeddings import EmbeddingGenerator
|
||||
from src.vector_store import FAISSVectorStore
|
||||
from src.retriever import DocumentRetriever
|
||||
from src.generator import AnswerGenerator
|
||||
|
||||
|
||||
class RAGSystem:
|
||||
"""Main RAG system orchestrator."""
|
||||
|
||||
def __init__(self, config: RAGConfig):
|
||||
self.config = config
|
||||
self.config.ensure_directories()
|
||||
|
||||
# Initialize components
|
||||
self.embedding_generator = EmbeddingGenerator(
|
||||
api_key=config.openai_api_key,
|
||||
model=config.embedding_model,
|
||||
batch_size=config.embedding_batch_size
|
||||
)
|
||||
|
||||
# Set embedding dimension based on model
|
||||
embedding_dim = ModelConfig.get_embedding_dimension(config.embedding_model)
|
||||
self.vector_store = FAISSVectorStore(dimension=embedding_dim)
|
||||
|
||||
self.retriever = DocumentRetriever(self.embedding_generator, self.vector_store)
|
||||
|
||||
self.generator = AnswerGenerator(
|
||||
api_key=config.openai_api_key,
|
||||
model=config.chat_model,
|
||||
temperature=config.temperature
|
||||
)
|
||||
|
||||
def build_index(self, force_rebuild: bool = False):
|
||||
"""Build or load the vector index."""
|
||||
index_path = self.config.get_index_path()
|
||||
metadata_path = self.config.get_metadata_path()
|
||||
|
||||
# Check if index exists and we don't need to rebuild
|
||||
if not force_rebuild and self.vector_store.exists(index_path, metadata_path):
|
||||
print("Loading existing index...")
|
||||
try:
|
||||
self.vector_store.load(index_path, metadata_path)
|
||||
print(f"Loaded index with {len(self.vector_store.texts)} chunks")
|
||||
return
|
||||
except Exception as e:
|
||||
print(f"Error loading index: {e}")
|
||||
print("Rebuilding index...")
|
||||
|
||||
# Load documents
|
||||
print(f"Loading documents from {self.config.documents_dir}...")
|
||||
loader = DocumentLoader(self.config.documents_dir)
|
||||
try:
|
||||
documents = loader.load_documents()
|
||||
except FileNotFoundError as e:
|
||||
print(f"Error: {e}")
|
||||
return
|
||||
|
||||
if not documents:
|
||||
print("No .txt files found in documents directory")
|
||||
return
|
||||
|
||||
print(f"Loaded {len(documents)} documents")
|
||||
|
||||
# Chunk documents
|
||||
print("Chunking documents...")
|
||||
chunker = TextChunker(
|
||||
chunk_size=self.config.chunk_size,
|
||||
overlap=self.config.chunk_overlap,
|
||||
model=self.config.embedding_model
|
||||
)
|
||||
chunks = chunker.chunk_documents(documents)
|
||||
print(f"Created {len(chunks)} chunks")
|
||||
|
||||
if not chunks:
|
||||
print("No chunks created")
|
||||
return
|
||||
|
||||
# Generate embeddings
|
||||
print("Generating embeddings...")
|
||||
texts = [chunk[0] for chunk in chunks]
|
||||
sources = [chunk[1] for chunk in chunks]
|
||||
|
||||
try:
|
||||
embeddings = self.embedding_generator.generate_embeddings(texts)
|
||||
except Exception as e:
|
||||
print(f"Error generating embeddings: {e}")
|
||||
return
|
||||
|
||||
# Create index
|
||||
print("Creating vector index...")
|
||||
self.vector_store.create_index(embeddings, texts, sources)
|
||||
|
||||
# Save index
|
||||
print("Saving index...")
|
||||
try:
|
||||
self.vector_store.save(index_path, metadata_path)
|
||||
print("Index saved successfully")
|
||||
except Exception as e:
|
||||
print(f"Error saving index: {e}")
|
||||
|
||||
def query(self, question: str) -> str:
|
||||
"""Process a single query."""
|
||||
if not question.strip():
|
||||
return "Please provide a valid question."
|
||||
|
||||
# Retrieve relevant chunks
|
||||
try:
|
||||
chunks = self.retriever.retrieve_with_threshold(
|
||||
question,
|
||||
top_k=self.config.top_k,
|
||||
threshold=self.config.similarity_threshold
|
||||
)
|
||||
except Exception as e:
|
||||
return f"Error during retrieval: {e}"
|
||||
|
||||
if not chunks:
|
||||
return "I couldn't find relevant information to answer your question."
|
||||
|
||||
# Generate answer
|
||||
try:
|
||||
answer = self.generator.generate_answer(question, chunks)
|
||||
return answer
|
||||
except Exception as e:
|
||||
return f"Error generating answer: {e}"
|
||||
|
||||
def interactive_mode(self):
|
||||
"""Run interactive query mode."""
|
||||
print("\n=== Interactive RAG System ===")
|
||||
print("Ask questions about your documents. Type 'quit' or 'exit' to stop.\n")
|
||||
|
||||
while True:
|
||||
try:
|
||||
question = input("Question: ").strip()
|
||||
|
||||
if question.lower() in ['quit', 'exit', 'q']:
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
if not question:
|
||||
continue
|
||||
|
||||
print("\nThinking...")
|
||||
answer = self.query(question)
|
||||
print(f"\nAnswer: {answer}\n")
|
||||
print("-" * 50)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main CLI entry point."""
|
||||
load_dotenv()
|
||||
|
||||
parser = argparse.ArgumentParser(description="RAG (Retrieval-Augmented Generation) System")
|
||||
parser.add_argument("--docs-dir", default="documents", help="Directory containing .txt documents")
|
||||
parser.add_argument("--index-dir", default="index", help="Directory for storing vector index")
|
||||
parser.add_argument("--question", "-q", help="Single question to ask")
|
||||
parser.add_argument("--interactive", "-i", action="store_true", help="Run in interactive mode")
|
||||
parser.add_argument("--rebuild-index", action="store_true", help="Force rebuild of vector index")
|
||||
parser.add_argument("--embedding-model", default="text-embedding-3-small",
|
||||
choices=list(ModelConfig.EMBEDDING_MODELS.keys()),
|
||||
help="OpenAI embedding model to use")
|
||||
parser.add_argument("--chat-model", default="gpt-3.5-turbo",
|
||||
choices=ModelConfig.CHAT_MODELS,
|
||||
help="OpenAI chat model to use")
|
||||
parser.add_argument("--top-k", type=int, default=5, help="Number of chunks to retrieve")
|
||||
parser.add_argument("--chunk-size", type=int, default=1000, help="Chunk size in tokens")
|
||||
parser.add_argument("--chunk-overlap", type=int, default=200, help="Chunk overlap in tokens")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get OpenAI API key
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OPENAI_API_KEY environment variable not set")
|
||||
print("Please set your OpenAI API key in the .env file or as an environment variable")
|
||||
sys.exit(1)
|
||||
|
||||
# Validate models
|
||||
if not ModelConfig.validate_models(args.embedding_model, args.chat_model):
|
||||
print("Error: Invalid model selection")
|
||||
sys.exit(1)
|
||||
|
||||
# Create configuration
|
||||
config = RAGConfig(
|
||||
openai_api_key=api_key,
|
||||
embedding_model=args.embedding_model,
|
||||
chat_model=args.chat_model,
|
||||
documents_dir=args.docs_dir,
|
||||
index_dir=args.index_dir,
|
||||
top_k=args.top_k,
|
||||
chunk_size=args.chunk_size,
|
||||
chunk_overlap=args.chunk_overlap
|
||||
)
|
||||
|
||||
# Initialize RAG system
|
||||
rag_system = RAGSystem(config)
|
||||
|
||||
# Build or load index
|
||||
rag_system.build_index(force_rebuild=args.rebuild_index)
|
||||
|
||||
# Process query or run interactive mode
|
||||
if args.question:
|
||||
answer = rag_system.query(args.question)
|
||||
print(f"Question: {args.question}")
|
||||
print(f"Answer: {answer}")
|
||||
elif args.interactive:
|
||||
rag_system.interactive_mode()
|
||||
else:
|
||||
print("Please provide either --question or --interactive flag")
|
||||
print("Use --help for more information")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
openai>=1.0.0
|
||||
faiss-cpu>=1.7.0
|
||||
tiktoken>=0.5.0
|
||||
numpy>=1.21.0
|
||||
python-dotenv>=1.0.0
|
||||
@@ -0,0 +1,72 @@
|
||||
import tiktoken
|
||||
from typing import List, Tuple
|
||||
|
||||
|
||||
class TextChunker:
|
||||
"""Splits text into overlapping chunks using tiktoken tokenization."""
|
||||
|
||||
def __init__(self, chunk_size: int = 1000, overlap: int = 200, model: str = "text-embedding-3-small"):
|
||||
self.chunk_size = chunk_size
|
||||
self.overlap = overlap
|
||||
self.encoding = tiktoken.encoding_for_model(model)
|
||||
|
||||
def chunk_text(self, text: str, source: str = "") -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Split text into overlapping chunks.
|
||||
|
||||
Args:
|
||||
text: Text content to chunk
|
||||
source: Source identifier (e.g., filename)
|
||||
|
||||
Returns:
|
||||
List of tuples containing (chunk_text, source)
|
||||
"""
|
||||
if not text.strip():
|
||||
return []
|
||||
|
||||
# Tokenize the text
|
||||
tokens = self.encoding.encode(text)
|
||||
|
||||
if len(tokens) <= self.chunk_size:
|
||||
return [(text, source)]
|
||||
|
||||
chunks = []
|
||||
start = 0
|
||||
|
||||
while start < len(tokens):
|
||||
# Determine end position
|
||||
end = start + self.chunk_size
|
||||
|
||||
# Extract chunk tokens
|
||||
chunk_tokens = tokens[start:end]
|
||||
|
||||
# Decode back to text
|
||||
chunk_text = self.encoding.decode(chunk_tokens)
|
||||
|
||||
chunks.append((chunk_text, source))
|
||||
|
||||
# Move start position with overlap
|
||||
if end >= len(tokens):
|
||||
break
|
||||
|
||||
start = end - self.overlap
|
||||
|
||||
return chunks
|
||||
|
||||
def chunk_documents(self, documents: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Chunk multiple documents.
|
||||
|
||||
Args:
|
||||
documents: List of (filename, content) tuples
|
||||
|
||||
Returns:
|
||||
List of (chunk_text, source) tuples
|
||||
"""
|
||||
all_chunks = []
|
||||
|
||||
for filename, content in documents:
|
||||
chunks = self.chunk_text(content, filename)
|
||||
all_chunks.extend(chunks)
|
||||
|
||||
return all_chunks
|
||||
@@ -0,0 +1,41 @@
|
||||
import os
|
||||
from typing import List, Tuple
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class DocumentLoader:
|
||||
"""Loads text documents from a directory."""
|
||||
|
||||
def __init__(self, documents_dir: str):
|
||||
self.documents_dir = Path(documents_dir)
|
||||
|
||||
def load_documents(self) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Load all .txt files from the documents directory.
|
||||
|
||||
Returns:
|
||||
List of tuples containing (filename, content)
|
||||
"""
|
||||
documents = []
|
||||
|
||||
if not self.documents_dir.exists():
|
||||
raise FileNotFoundError(f"Documents directory not found: {self.documents_dir}")
|
||||
|
||||
for file_path in self.documents_dir.glob("*.txt"):
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
documents.append((file_path.name, content))
|
||||
except UnicodeDecodeError:
|
||||
try:
|
||||
with open(file_path, 'r', encoding='latin-1') as file:
|
||||
content = file.read()
|
||||
documents.append((file_path.name, content))
|
||||
except Exception as e:
|
||||
print(f"Error reading file {file_path}: {e}")
|
||||
continue
|
||||
except Exception as e:
|
||||
print(f"Error reading file {file_path}: {e}")
|
||||
continue
|
||||
|
||||
return documents
|
||||
@@ -0,0 +1,68 @@
|
||||
import asyncio
|
||||
import numpy as np
|
||||
from typing import List
|
||||
from openai import OpenAI
|
||||
import time
|
||||
|
||||
|
||||
class EmbeddingGenerator:
|
||||
"""Wraps OpenAI embedding API calls with batching support."""
|
||||
|
||||
def __init__(self, api_key: str, model: str = "text-embedding-3-small", batch_size: int = 100):
|
||||
self.client = OpenAI(api_key=api_key)
|
||||
self.model = model
|
||||
self.batch_size = batch_size
|
||||
|
||||
def generate_embeddings(self, texts: List[str]) -> np.ndarray:
|
||||
"""
|
||||
Generate embeddings for a list of texts with batching.
|
||||
|
||||
Args:
|
||||
texts: List of text strings
|
||||
|
||||
Returns:
|
||||
numpy array of embeddings
|
||||
"""
|
||||
if not texts:
|
||||
return np.array([])
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i:i + self.batch_size]
|
||||
batch_embeddings = self._generate_batch_embeddings(batch)
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
# Rate limiting - sleep briefly between batches
|
||||
if i + self.batch_size < len(texts):
|
||||
time.sleep(0.1)
|
||||
|
||||
return np.array(all_embeddings)
|
||||
|
||||
def _generate_batch_embeddings(self, texts: List[str]) -> List[List[float]]:
|
||||
"""Generate embeddings for a single batch."""
|
||||
max_retries = 3
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = self.client.embeddings.create(
|
||||
input=texts,
|
||||
model=self.model
|
||||
)
|
||||
return [embedding.embedding for embedding in response.data]
|
||||
|
||||
except Exception as e:
|
||||
if attempt == max_retries - 1:
|
||||
raise e
|
||||
|
||||
wait_time = 2 ** attempt
|
||||
print(f"Error generating embeddings (attempt {attempt + 1}): {e}")
|
||||
print(f"Retrying in {wait_time} seconds...")
|
||||
time.sleep(wait_time)
|
||||
|
||||
return []
|
||||
|
||||
def generate_single_embedding(self, text: str) -> np.ndarray:
|
||||
"""Generate embedding for a single text."""
|
||||
embeddings = self.generate_embeddings([text])
|
||||
return embeddings[0] if len(embeddings) > 0 else np.array([])
|
||||
@@ -0,0 +1,130 @@
|
||||
from typing import List, Tuple
|
||||
from openai import OpenAI
|
||||
|
||||
|
||||
class AnswerGenerator:
|
||||
"""Generates answers using OpenAI GPT with retrieved context."""
|
||||
|
||||
def __init__(self, api_key: str, model: str = "gpt-3.5-turbo", temperature: float = 0.1):
|
||||
self.client = OpenAI(api_key=api_key)
|
||||
self.model = model
|
||||
self.temperature = temperature
|
||||
|
||||
def generate_answer(self, query: str, context_chunks: List[Tuple[str, str, float]]) -> str:
|
||||
"""
|
||||
Generate an answer based on query and retrieved context.
|
||||
|
||||
Args:
|
||||
query: User question
|
||||
context_chunks: List of (text, source, score) tuples
|
||||
|
||||
Returns:
|
||||
Generated answer
|
||||
"""
|
||||
if not context_chunks:
|
||||
return "I couldn't find relevant information to answer your question."
|
||||
|
||||
# Prepare context
|
||||
context_parts = []
|
||||
sources = set()
|
||||
|
||||
for text, source, score in context_chunks:
|
||||
context_parts.append(f"From {source}: {text}")
|
||||
sources.add(source)
|
||||
|
||||
context = "\n\n".join(context_parts)
|
||||
|
||||
# Create prompt
|
||||
prompt = f"""Based on the following context, please answer the user's question. If the context doesn't contain enough information to answer the question, say so clearly.
|
||||
|
||||
Context:
|
||||
{context}
|
||||
|
||||
Question: {query}
|
||||
|
||||
Answer:"""
|
||||
|
||||
try:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant that answers questions based on provided context. Always be accurate and cite your sources when possible."
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt
|
||||
}
|
||||
],
|
||||
temperature=self.temperature,
|
||||
max_tokens=1000
|
||||
)
|
||||
|
||||
answer = response.choices[0].message.content.strip()
|
||||
|
||||
# Add sources information
|
||||
if sources:
|
||||
source_list = ", ".join(sorted(sources))
|
||||
answer += f"\n\nSources: {source_list}"
|
||||
|
||||
return answer
|
||||
|
||||
except Exception as e:
|
||||
return f"Error generating answer: {str(e)}"
|
||||
|
||||
def generate_streaming_answer(self, query: str, context_chunks: List[Tuple[str, str, float]]):
|
||||
"""Generate answer with streaming response."""
|
||||
if not context_chunks:
|
||||
yield "I couldn't find relevant information to answer your question."
|
||||
return
|
||||
|
||||
# Prepare context
|
||||
context_parts = []
|
||||
sources = set()
|
||||
|
||||
for text, source, score in context_chunks:
|
||||
context_parts.append(f"From {source}: {text}")
|
||||
sources.add(source)
|
||||
|
||||
context = "\n\n".join(context_parts)
|
||||
|
||||
# Create prompt
|
||||
prompt = f"""Based on the following context, please answer the user's question. If the context doesn't contain enough information to answer the question, say so clearly.
|
||||
|
||||
Context:
|
||||
{context}
|
||||
|
||||
Question: {query}
|
||||
|
||||
Answer:"""
|
||||
|
||||
try:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant that answers questions based on provided context. Always be accurate and cite your sources when possible."
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt
|
||||
}
|
||||
],
|
||||
temperature=self.temperature,
|
||||
max_tokens=1000,
|
||||
stream=True
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.content is not None:
|
||||
yield chunk.choices[0].delta.content
|
||||
|
||||
# Add sources information
|
||||
if sources:
|
||||
source_list = ", ".join(sorted(sources))
|
||||
yield f"\n\nSources: {source_list}"
|
||||
|
||||
except Exception as e:
|
||||
yield f"Error generating answer: {str(e)}"
|
||||
@@ -0,0 +1,60 @@
|
||||
from typing import List, Tuple
|
||||
import numpy as np
|
||||
from .embeddings import EmbeddingGenerator
|
||||
from .vector_store import FAISSVectorStore
|
||||
|
||||
|
||||
class DocumentRetriever:
|
||||
"""Combines embedding generation and vector search."""
|
||||
|
||||
def __init__(self, embedding_generator: EmbeddingGenerator, vector_store: FAISSVectorStore):
|
||||
self.embedding_generator = embedding_generator
|
||||
self.vector_store = vector_store
|
||||
|
||||
def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[str, str, float]]:
|
||||
"""
|
||||
Retrieve relevant document chunks for a query.
|
||||
|
||||
Args:
|
||||
query: User question
|
||||
top_k: Number of chunks to retrieve
|
||||
|
||||
Returns:
|
||||
List of tuples containing (text, source, score)
|
||||
"""
|
||||
if not query.strip():
|
||||
return []
|
||||
|
||||
# Generate query embedding
|
||||
query_embedding = self.embedding_generator.generate_single_embedding(query)
|
||||
|
||||
if len(query_embedding) == 0:
|
||||
return []
|
||||
|
||||
# Search for similar chunks
|
||||
results = self.vector_store.search(query_embedding, top_k)
|
||||
|
||||
return results
|
||||
|
||||
def retrieve_with_threshold(self, query: str, top_k: int = 5, threshold: float = 0.7) -> List[Tuple[str, str, float]]:
|
||||
"""
|
||||
Retrieve relevant chunks with a similarity threshold.
|
||||
|
||||
Args:
|
||||
query: User question
|
||||
top_k: Number of chunks to retrieve
|
||||
threshold: Minimum similarity score
|
||||
|
||||
Returns:
|
||||
List of tuples containing (text, source, score)
|
||||
"""
|
||||
results = self.retrieve(query, top_k)
|
||||
|
||||
# Filter by threshold
|
||||
filtered_results = [
|
||||
(text, source, score)
|
||||
for text, source, score in results
|
||||
if score >= threshold
|
||||
]
|
||||
|
||||
return filtered_results
|
||||
@@ -0,0 +1,125 @@
|
||||
import faiss
|
||||
import numpy as np
|
||||
import pickle
|
||||
from typing import List, Tuple, Optional
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class FAISSVectorStore:
|
||||
"""Manages FAISS index operations with persistence."""
|
||||
|
||||
def __init__(self, dimension: int = 1536):
|
||||
self.dimension = dimension
|
||||
self.index = None
|
||||
self.texts = []
|
||||
self.sources = []
|
||||
|
||||
def create_index(self, embeddings: np.ndarray, texts: List[str], sources: List[str]):
|
||||
"""
|
||||
Create a new FAISS index.
|
||||
|
||||
Args:
|
||||
embeddings: Array of embeddings
|
||||
texts: List of text chunks
|
||||
sources: List of source identifiers
|
||||
"""
|
||||
if len(embeddings) == 0:
|
||||
raise ValueError("Cannot create index with empty embeddings")
|
||||
|
||||
if embeddings.shape[1] != self.dimension:
|
||||
self.dimension = embeddings.shape[1]
|
||||
|
||||
# Create FAISS index
|
||||
self.index = faiss.IndexFlatIP(self.dimension) # Inner product for cosine similarity
|
||||
|
||||
# Normalize embeddings for cosine similarity
|
||||
normalized_embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||
|
||||
# Add embeddings to index
|
||||
self.index.add(normalized_embeddings.astype('float32'))
|
||||
|
||||
# Store metadata
|
||||
self.texts = texts.copy()
|
||||
self.sources = sources.copy()
|
||||
|
||||
def add_vectors(self, embeddings: np.ndarray, texts: List[str], sources: List[str]):
|
||||
"""Add new vectors to existing index."""
|
||||
if self.index is None:
|
||||
self.create_index(embeddings, texts, sources)
|
||||
return
|
||||
|
||||
# Normalize embeddings
|
||||
normalized_embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||
|
||||
# Add to index
|
||||
self.index.add(normalized_embeddings.astype('float32'))
|
||||
|
||||
# Update metadata
|
||||
self.texts.extend(texts)
|
||||
self.sources.extend(sources)
|
||||
|
||||
def search(self, query_embedding: np.ndarray, top_k: int = 5) -> List[Tuple[str, str, float]]:
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
query_embedding: Query vector
|
||||
top_k: Number of results to return
|
||||
|
||||
Returns:
|
||||
List of tuples containing (text, source, score)
|
||||
"""
|
||||
if self.index is None or self.index.ntotal == 0:
|
||||
return []
|
||||
|
||||
# Normalize query embedding
|
||||
query_embedding = query_embedding / np.linalg.norm(query_embedding)
|
||||
query_embedding = query_embedding.reshape(1, -1).astype('float32')
|
||||
|
||||
# Search
|
||||
scores, indices = self.index.search(query_embedding, min(top_k, self.index.ntotal))
|
||||
|
||||
results = []
|
||||
for score, idx in zip(scores[0], indices[0]):
|
||||
if idx != -1 and idx < len(self.texts):
|
||||
results.append((self.texts[idx], self.sources[idx], float(score)))
|
||||
|
||||
return results
|
||||
|
||||
def save(self, index_path: str, metadata_path: str):
|
||||
"""Save index and metadata to disk."""
|
||||
if self.index is None:
|
||||
raise ValueError("No index to save")
|
||||
|
||||
# Save FAISS index
|
||||
faiss.write_index(self.index, index_path)
|
||||
|
||||
# Save metadata
|
||||
metadata = {
|
||||
'texts': self.texts,
|
||||
'sources': self.sources,
|
||||
'dimension': self.dimension
|
||||
}
|
||||
|
||||
with open(metadata_path, 'wb') as f:
|
||||
pickle.dump(metadata, f)
|
||||
|
||||
def load(self, index_path: str, metadata_path: str):
|
||||
"""Load index and metadata from disk."""
|
||||
if not Path(index_path).exists() or not Path(metadata_path).exists():
|
||||
raise FileNotFoundError("Index or metadata file not found")
|
||||
|
||||
# Load FAISS index
|
||||
self.index = faiss.read_index(index_path)
|
||||
|
||||
# Load metadata
|
||||
with open(metadata_path, 'rb') as f:
|
||||
metadata = pickle.load(f)
|
||||
|
||||
self.texts = metadata['texts']
|
||||
self.sources = metadata['sources']
|
||||
self.dimension = metadata['dimension']
|
||||
|
||||
def exists(self, index_path: str, metadata_path: str) -> bool:
|
||||
"""Check if saved index exists."""
|
||||
return Path(index_path).exists() and Path(metadata_path).exists()
|
||||
@@ -61,6 +61,113 @@ _SANITIZER = PromptSanitizer()
|
||||
|
||||
PATH_HINT_RE = re.compile(r"@(?P<path>[\w./-]+)")
|
||||
|
||||
# Regex for extracting files from markdown-formatted LLM output.
|
||||
# Supports many common header patterns:
|
||||
# ## filename.py ## `filename.py`
|
||||
# ## 1. `filename.py` ### path/to/file.py
|
||||
# ## **filename.py** # filename.py
|
||||
# **filename.py** **`filename.py`**
|
||||
_FILE_HEADER_RE = re.compile(
|
||||
r"^(?:#{1,3}\s+(?:\d+\.\s+)?)?[`*]*(?P<path>[^\s`*]+\.[a-zA-Z0-9]+)[`*]*\s*$",
|
||||
re.MULTILINE,
|
||||
)
|
||||
_FENCED_BLOCK_RE = re.compile(
|
||||
r"```[a-zA-Z]*\s*\n(?P<code>.*?)```",
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
def _extract_files_from_markdown(text: str) -> list[tuple[str, str]]:
|
||||
"""Extract (file_path, code_content) pairs from markdown LLM output.
|
||||
|
||||
Parses output formatted as markdown headers followed by fenced code blocks::
|
||||
|
||||
## filename.py
|
||||
```python
|
||||
code here
|
||||
```
|
||||
|
||||
Handles many header variants: ``#``, ``##``, ``###``, numbered (``1.``),
|
||||
backtick-wrapped, bold (``**``), and nested paths (``templates/base.html``).
|
||||
|
||||
Returns a list of (path, content) tuples. If fewer than 2 file headers
|
||||
are found, returns an empty list (caller falls back to single-file).
|
||||
"""
|
||||
files: list[tuple[str, str]] = []
|
||||
seen_paths: set[str] = set()
|
||||
|
||||
# Build a set of character ranges that are inside fenced code blocks
|
||||
# so we can skip header matches that fall within them.
|
||||
fenced_ranges: list[tuple[int, int]] = []
|
||||
for block in _FENCED_BLOCK_RE.finditer(text):
|
||||
fenced_ranges.append((block.start(), block.end()))
|
||||
|
||||
def _inside_fence(pos: int) -> bool:
|
||||
return any(start <= pos < end for start, end in fenced_ranges)
|
||||
|
||||
# Find all heading lines that look like file paths, excluding those
|
||||
# inside fenced code blocks (e.g. dependency lines like requests>=2.31.0)
|
||||
headers = [h for h in _FILE_HEADER_RE.finditer(text) if not _inside_fence(h.start())]
|
||||
if len(headers) < 2:
|
||||
# Need at least 2 file headers to confirm multi-file intent.
|
||||
# A single header could be a coincidental match in prose.
|
||||
return []
|
||||
|
||||
for i, header in enumerate(headers):
|
||||
file_path = header.group("path").strip("`*")
|
||||
|
||||
# Skip non-file headers (e.g., ## Features, ## Installation)
|
||||
if "/" not in file_path and "." not in file_path:
|
||||
continue
|
||||
|
||||
# Deduplicate
|
||||
if file_path in seen_paths:
|
||||
continue
|
||||
|
||||
# Get the text between this header and the next (or end of text)
|
||||
start = header.end()
|
||||
end = headers[i + 1].start() if i + 1 < len(headers) else len(text)
|
||||
section = text[start:end]
|
||||
|
||||
# Extract the first fenced code block in this section
|
||||
block_match = _FENCED_BLOCK_RE.search(section)
|
||||
if block_match:
|
||||
code = block_match.group("code").strip()
|
||||
if code:
|
||||
seen_paths.add(file_path)
|
||||
files.append((file_path, code))
|
||||
|
||||
return files
|
||||
|
||||
|
||||
def _extract_single_file(text: str) -> tuple[str, str]:
|
||||
"""Extract a single file from LLM output when multi-file extraction fails.
|
||||
|
||||
Tries to find a filename header (``**filename.py**`` or ``## filename.py``)
|
||||
followed by a fenced code block. If no header is found, extracts the first
|
||||
fenced block and uses ``generated.py`` as the path. If no fenced block
|
||||
exists either, returns the raw text as ``generated.py``.
|
||||
|
||||
Returns:
|
||||
(file_path, code_content) tuple.
|
||||
"""
|
||||
# Look for a single file header anywhere in the text
|
||||
header = _FILE_HEADER_RE.search(text)
|
||||
block = _FENCED_BLOCK_RE.search(text)
|
||||
|
||||
if header and block:
|
||||
file_path = header.group("path").strip("`*")
|
||||
code = block.group("code").strip()
|
||||
return (file_path, code)
|
||||
|
||||
if block:
|
||||
# Fenced block but no recognizable header
|
||||
code = block.group("code").strip()
|
||||
return ("generated.py", code)
|
||||
|
||||
# No fenced block at all — return raw text
|
||||
return ("generated.py", text.strip())
|
||||
|
||||
|
||||
class BoundedMemorySaver(MemorySaver):
|
||||
"""Memory saver that retains only the latest checkpoints per thread."""
|
||||
@@ -211,43 +318,72 @@ class PlanGenerationGraph:
|
||||
input_variables=["prompt", "context_summary"],
|
||||
template=(
|
||||
f"{_bi}\n\n"
|
||||
"You are a requirements analyst. "
|
||||
"Extract concise technical requirements from the request "
|
||||
"and context.\n\n"
|
||||
"You are a senior software architect. Analyze the request below "
|
||||
"and produce a structured implementation plan.\n\n"
|
||||
f"Request: {_bs}\n{{prompt}}\n{_be}\n"
|
||||
f"Context:\n{_bs}\n{{context_summary}}\n{_be}\n\n"
|
||||
"Return:\n"
|
||||
"- Key requirements\n"
|
||||
"- Files to modify or create\n"
|
||||
"- Needed dependencies\n"
|
||||
"- Risks or challenges\n"
|
||||
"Return a structured analysis with these sections:\n"
|
||||
"## Files to Create\n"
|
||||
"List each file with its purpose (e.g., `app.py` — Flask application entry point)\n\n"
|
||||
"## Dependencies\n"
|
||||
"List all third-party packages needed (e.g., flask>=2.0.0)\n\n"
|
||||
"## Key Requirements\n"
|
||||
"Numbered list of functional requirements extracted from the request\n\n"
|
||||
"## Architecture Notes\n"
|
||||
"Brief notes on structure, patterns, and how files relate to each other\n"
|
||||
),
|
||||
)
|
||||
|
||||
# Code generation prompt
|
||||
# Code generation prompt — structured output format
|
||||
self.generate_prompt: Any = PromptTemplate(
|
||||
input_variables=["requirements", "context_summary"],
|
||||
input_variables=["requirements", "context_summary", "original_prompt"],
|
||||
template=(
|
||||
f"{_bi}\n\n"
|
||||
"You are a code generator. "
|
||||
"Produce production-ready code that satisfies the requirements using "
|
||||
"the provided context.\n\n"
|
||||
f"Requirements:\n{_bs}\n{{requirements}}\n{_be}\n\n"
|
||||
"You are an expert software engineer. Generate complete, "
|
||||
"production-ready code based on the requirements below.\n\n"
|
||||
f"Original Request & Definition of Done:\n{_bs}\n{{original_prompt}}\n{_be}\n\n"
|
||||
f"Architecture Analysis:\n{_bs}\n{{requirements}}\n{_be}\n\n"
|
||||
f"Context:\n{_bs}\n{{context_summary}}\n{_be}\n\n"
|
||||
"Return concise, documented code that follows best practices.\n"
|
||||
"## Output Format Rules (CRITICAL — follow exactly)\n\n"
|
||||
"You MUST output ALL files listed in the Definition of Done above. "
|
||||
"Each file must be a separate code block with its own header.\n\n"
|
||||
"Output each file using this exact format:\n\n"
|
||||
"**filename.ext**\n"
|
||||
"```language\n"
|
||||
"file contents here\n"
|
||||
"```\n\n"
|
||||
"Rules:\n"
|
||||
"- Output EVERY file mentioned in the Definition of Done — "
|
||||
"missing files is a failure\n"
|
||||
"- Use the **filename.ext** header on its own line before each "
|
||||
"fenced code block\n"
|
||||
"- Use correct file paths including subdirectories "
|
||||
"(e.g., **src/module.py**, **templates/index.html**)\n"
|
||||
"- Each file must be complete and runnable — no placeholders, "
|
||||
"no TODOs, no `...` or `pass`\n"
|
||||
"- Include all imports at the top of each file\n"
|
||||
"- Include proper error handling\n"
|
||||
"- Do NOT include explanatory prose between files — only the "
|
||||
"file header and code block\n"
|
||||
),
|
||||
)
|
||||
|
||||
# Validation prompt
|
||||
# Validation prompt — structured review
|
||||
self.validate_prompt: Any = PromptTemplate(
|
||||
input_variables=["generated_code"],
|
||||
template=(
|
||||
f"{_bi}\n\n"
|
||||
"You are a reviewer. "
|
||||
"Validate the generated code for correctness and safety.\n\n"
|
||||
"You are a code reviewer. Review the generated code below "
|
||||
"for correctness.\n\n"
|
||||
f"Code:\n{_bs}\n{{generated_code}}\n{_be}\n\n"
|
||||
"Check syntax, logic, best practices, security, and performance. "
|
||||
"Respond with PASS or FAIL and list issues found.\n"
|
||||
"Check for:\n"
|
||||
"1. Missing imports\n"
|
||||
"2. Syntax errors\n"
|
||||
"3. Undefined variables or functions\n"
|
||||
"4. Security issues (SQL injection, XSS, hardcoded secrets)\n"
|
||||
"5. Missing error handling for I/O operations\n\n"
|
||||
"Start your response with exactly PASS or FAIL on the first line.\n"
|
||||
"If FAIL, list each issue on a separate line starting with '- '.\n"
|
||||
),
|
||||
)
|
||||
|
||||
@@ -388,76 +524,72 @@ class PlanGenerationGraph:
|
||||
)
|
||||
|
||||
try:
|
||||
# Generate code
|
||||
# Generate code — pass original prompt so LLM sees the full DoD
|
||||
result = chain.invoke(
|
||||
{
|
||||
"requirements": requirements.get("description", ""),
|
||||
"context_summary": context_summary,
|
||||
"original_prompt": state.get("prompt", ""),
|
||||
}
|
||||
)
|
||||
generated_code = str(result)
|
||||
|
||||
# Determine file path and operation
|
||||
operation = requirements.get("operation", "create")
|
||||
contexts = state.get("contexts", [])
|
||||
explicit_path = self._extract_path_from_prompt(state.get("prompt", ""))
|
||||
|
||||
if operation == "modify" and contexts:
|
||||
matched_context = self._find_matching_context(contexts, explicit_path)
|
||||
target_context = matched_context or contexts[0]
|
||||
file_path = target_context.path
|
||||
operation_type = OperationType.MODIFY
|
||||
original_content = target_context.content
|
||||
elif explicit_path:
|
||||
file_path = explicit_path
|
||||
operation_type = (
|
||||
OperationType.MODIFY if contexts else OperationType.CREATE
|
||||
)
|
||||
original_content = None
|
||||
|
||||
existing_path = Path(file_path)
|
||||
if existing_path.exists():
|
||||
if existing_path.is_dir():
|
||||
file_path = str(existing_path / "generated.py")
|
||||
operation_type = OperationType.CREATE
|
||||
original_content = None
|
||||
else:
|
||||
operation_type = OperationType.MODIFY
|
||||
try:
|
||||
original_content = existing_path.read_text()
|
||||
except Exception: # pragma: no cover - best effort
|
||||
original_content = None
|
||||
elif existing_path.suffix == "":
|
||||
file_path = str(existing_path / "generated.py")
|
||||
operation_type = OperationType.CREATE
|
||||
original_content = None
|
||||
else:
|
||||
prompt_lower = state["prompt"].lower()
|
||||
if "test" in prompt_lower:
|
||||
file_path = "test_generated.py"
|
||||
elif "error" in prompt_lower or "exception" in prompt_lower:
|
||||
file_path = "error_handler.py"
|
||||
else:
|
||||
file_path = "generated.py"
|
||||
operation_type = OperationType.CREATE
|
||||
original_content = None
|
||||
|
||||
# Create change object
|
||||
plan_id = state["plan"].id if state["plan"].id else 0
|
||||
|
||||
changes = [
|
||||
Change(
|
||||
id=None,
|
||||
plan_id=plan_id,
|
||||
file_path=file_path,
|
||||
operation=operation_type,
|
||||
original_content=original_content,
|
||||
new_content=generated_code,
|
||||
applied=False,
|
||||
applied_at=None,
|
||||
new_path=None,
|
||||
)
|
||||
]
|
||||
# Try to extract multiple files from markdown-formatted output
|
||||
extracted_files = _extract_files_from_markdown(generated_code)
|
||||
|
||||
if extracted_files:
|
||||
# Multi-file output: create one Change per extracted file
|
||||
changes = []
|
||||
for extracted_path, extracted_code in extracted_files:
|
||||
changes.append(
|
||||
Change(
|
||||
id=None,
|
||||
plan_id=plan_id,
|
||||
file_path=extracted_path,
|
||||
operation=OperationType.CREATE,
|
||||
original_content=None,
|
||||
new_content=extracted_code,
|
||||
applied=False,
|
||||
applied_at=None,
|
||||
new_path=None,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# Single-file fallback: try to extract one file from a
|
||||
# single fenced block with a preceding filename header
|
||||
file_path, code_content = _extract_single_file(generated_code)
|
||||
|
||||
# Determine operation type from context
|
||||
contexts = state.get("contexts", [])
|
||||
operation = requirements.get("operation", "create")
|
||||
original_content: str | None = None
|
||||
operation_type = OperationType.CREATE
|
||||
|
||||
if operation == "modify" and contexts:
|
||||
explicit_path = self._extract_path_from_prompt(
|
||||
state.get("prompt", "")
|
||||
)
|
||||
matched = self._find_matching_context(contexts, explicit_path)
|
||||
target = matched or contexts[0]
|
||||
file_path = target.path
|
||||
operation_type = OperationType.MODIFY
|
||||
original_content = target.content
|
||||
|
||||
changes = [
|
||||
Change(
|
||||
id=None,
|
||||
plan_id=plan_id,
|
||||
file_path=file_path,
|
||||
operation=operation_type,
|
||||
original_content=original_content,
|
||||
new_content=code_content,
|
||||
applied=False,
|
||||
applied_at=None,
|
||||
new_path=None,
|
||||
)
|
||||
]
|
||||
|
||||
return {
|
||||
"generated_changes": changes,
|
||||
@@ -493,6 +625,10 @@ class PlanGenerationGraph:
|
||||
def _validate(self, state: PlanGenerationState) -> dict[str, Any]:
|
||||
"""Validate generated code changes.
|
||||
|
||||
Performs two levels of validation:
|
||||
1. Local syntax check (``compile()``) for Python files
|
||||
2. LLM-based review for logic, security, and completeness
|
||||
|
||||
Args:
|
||||
state: Current workflow state
|
||||
|
||||
@@ -509,26 +645,44 @@ class PlanGenerationGraph:
|
||||
},
|
||||
}
|
||||
|
||||
# Create validation chain
|
||||
# --- Phase 1: local syntax check for Python files ---
|
||||
syntax_errors: list[str] = []
|
||||
for change in changes:
|
||||
if not change.file_path.endswith(".py"):
|
||||
continue
|
||||
code = change.new_content or ""
|
||||
try:
|
||||
compile(code, change.file_path, "exec")
|
||||
except SyntaxError as exc:
|
||||
line_info = f" (line {exc.lineno})" if exc.lineno else ""
|
||||
syntax_errors.append(
|
||||
f"{change.file_path}{line_info}: {exc.msg}"
|
||||
)
|
||||
|
||||
if syntax_errors:
|
||||
return {
|
||||
"validation_result": {
|
||||
"status": "FAIL",
|
||||
"message": "Syntax errors:\n" + "\n".join(syntax_errors),
|
||||
},
|
||||
}
|
||||
|
||||
# --- Phase 2: LLM-based review ---
|
||||
chain = self._chain_with_retry(
|
||||
self.validate_prompt | self.llm | StrOutputParser()
|
||||
)
|
||||
|
||||
try:
|
||||
# Validate all changes
|
||||
all_code = "\n\n".join(
|
||||
f"File: {change.file_path}\n{change.new_content}" for change in changes
|
||||
)
|
||||
|
||||
result = chain.invoke(
|
||||
{
|
||||
"generated_code": all_code,
|
||||
}
|
||||
)
|
||||
validation = str(result)
|
||||
result = chain.invoke({"generated_code": all_code})
|
||||
validation = str(result).strip()
|
||||
|
||||
# Simple validation check (in real implementation, parse the LLM response)
|
||||
is_valid = "PASS" in validation.upper() or len(all_code) > 10
|
||||
# Parse first line for PASS/FAIL verdict
|
||||
first_line = validation.split("\n", 1)[0].strip().upper()
|
||||
is_valid = first_line.startswith("PASS")
|
||||
|
||||
return {
|
||||
"validation_result": {
|
||||
@@ -538,10 +692,11 @@ class PlanGenerationGraph:
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
# If LLM review fails, still pass if syntax was clean
|
||||
return {
|
||||
"validation_result": {
|
||||
"status": "FAIL",
|
||||
"message": f"Validation failed: {e!s}",
|
||||
"status": "PASS",
|
||||
"message": f"LLM review unavailable ({e!s}); syntax check passed.",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -280,6 +280,23 @@ def _build_resource_file_watcher(
|
||||
)
|
||||
|
||||
|
||||
def _build_changeset_store(database_url: str) -> object:
|
||||
"""Build a ChangeSetStore for recording changeset entries.
|
||||
|
||||
Currently returns ``InMemoryChangeSetStore`` (matching the M1 spec).
|
||||
Cross-process persistence of generated file content is handled by the
|
||||
sandbox directory, not the changeset store. The store provides
|
||||
in-process audit tracking during ``run_execute``.
|
||||
|
||||
A future milestone will upgrade to ``SqliteChangeSetStore`` once
|
||||
session management is unified with the UoW to avoid SQLite lock
|
||||
contention from concurrent connections.
|
||||
"""
|
||||
from cleveragents.domain.models.core.change import InMemoryChangeSetStore
|
||||
|
||||
return InMemoryChangeSetStore()
|
||||
|
||||
|
||||
def _build_trace_service(
|
||||
database_url: str,
|
||||
settings: Settings | None = None,
|
||||
@@ -429,6 +446,12 @@ class Container(containers.DeclarativeContainer):
|
||||
event_bus=event_bus,
|
||||
)
|
||||
|
||||
# Changeset Store - Singleton (SQLite-backed changeset persistence)
|
||||
changeset_store = providers.Singleton(
|
||||
_build_changeset_store,
|
||||
database_url=database_url,
|
||||
)
|
||||
|
||||
# Plan Lifecycle Service - Factory (v3 four-phase lifecycle)
|
||||
plan_lifecycle_service = providers.Factory(
|
||||
PlanLifecycleService,
|
||||
@@ -436,6 +459,10 @@ class Container(containers.DeclarativeContainer):
|
||||
unit_of_work=unit_of_work,
|
||||
decision_service=decision_service,
|
||||
event_bus=event_bus,
|
||||
ai_provider=ai_provider,
|
||||
provider_registry=provider_registry,
|
||||
actor_service=actor_service,
|
||||
changeset_store=changeset_store,
|
||||
)
|
||||
|
||||
# Checkpoint Service - database-backed via CheckpointRepository
|
||||
|
||||
@@ -88,14 +88,20 @@ from cleveragents.infrastructure.events.models import DomainEvent
|
||||
from cleveragents.infrastructure.events.types import EventType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from cleveragents.application.services.actor_service import ActorService
|
||||
from cleveragents.application.services.async_worker import InMemoryJobStore
|
||||
from cleveragents.application.services.decision_service import DecisionService
|
||||
from cleveragents.application.services.error_pattern_service import (
|
||||
ErrorPatternService,
|
||||
)
|
||||
from cleveragents.config.settings import Settings
|
||||
from cleveragents.domain.models.core.change import ChangeSetStore
|
||||
from cleveragents.domain.providers.ai_provider import AIProviderInterface
|
||||
from cleveragents.infrastructure.database.unit_of_work import UnitOfWork
|
||||
from cleveragents.infrastructure.events.protocol import EventBus
|
||||
from cleveragents.providers.registry import ProviderRegistry
|
||||
|
||||
logger = structlog.get_logger(__name__)
|
||||
|
||||
@@ -163,6 +169,10 @@ class PlanLifecycleService:
|
||||
event_bus: EventBus | None = None,
|
||||
job_store: InMemoryJobStore | None = None,
|
||||
error_pattern_service: ErrorPatternService | None = None,
|
||||
ai_provider: AIProviderInterface | None = None,
|
||||
provider_registry: ProviderRegistry | None = None,
|
||||
actor_service: ActorService | None = None,
|
||||
changeset_store: ChangeSetStore | None = None,
|
||||
):
|
||||
"""Initialize the plan lifecycle service.
|
||||
|
||||
@@ -191,6 +201,18 @@ class PlanLifecycleService:
|
||||
consulted before the Execute phase to inject preventive
|
||||
guidance. When ``None``, predictive prevention is
|
||||
silently skipped.
|
||||
ai_provider: Optional AI provider for executing plans with
|
||||
real LLM calls. When ``None``, the provider_registry
|
||||
is used to create one on demand.
|
||||
provider_registry: Optional provider registry for creating
|
||||
AI providers. Used when ai_provider is not provided.
|
||||
actor_service: Optional :class:`ActorService` for resolving
|
||||
execution actors by name. When provided, actors are
|
||||
looked up via the database, extracting pre-stored
|
||||
``provider`` and ``model`` fields.
|
||||
changeset_store: Optional :class:`ChangeSetStore` for
|
||||
recording changeset entries during execution. Provides
|
||||
auditable tracking of generated file changes.
|
||||
"""
|
||||
self.settings = settings
|
||||
self.unit_of_work = unit_of_work
|
||||
@@ -198,6 +220,10 @@ class PlanLifecycleService:
|
||||
self.event_bus = event_bus
|
||||
self._job_store = job_store
|
||||
self.error_pattern_service = error_pattern_service
|
||||
self._ai_provider = ai_provider
|
||||
self._provider_registry = provider_registry
|
||||
self._actor_service = actor_service
|
||||
self._changeset_store = changeset_store
|
||||
self._logger = logger.bind(service="plan_lifecycle")
|
||||
self.preflight_guardrail = PlanPreflightGuardrail()
|
||||
|
||||
@@ -553,7 +579,30 @@ class PlanLifecycleService:
|
||||
Returns:
|
||||
List of matching Actions
|
||||
"""
|
||||
actions = list(self._actions.values())
|
||||
# Start with in-memory actions
|
||||
seen_names: set[str] = set()
|
||||
actions: list[Action] = []
|
||||
|
||||
for action in self._actions.values():
|
||||
seen_names.add(str(action.namespaced_name))
|
||||
actions.append(action)
|
||||
|
||||
# Merge from persistence layer
|
||||
if self._persisted and self.unit_of_work is not None:
|
||||
try:
|
||||
with self.unit_of_work.transaction() as ctx:
|
||||
persisted = ctx.actions.list_available(namespace=namespace)
|
||||
for action in persisted:
|
||||
name_str = str(action.namespaced_name)
|
||||
if name_str not in seen_names:
|
||||
seen_names.add(name_str)
|
||||
actions.append(action)
|
||||
self._actions[name_str] = action
|
||||
except Exception:
|
||||
self._logger.warning(
|
||||
"list_actions_persistence_fallback",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
if namespace:
|
||||
actions = [a for a in actions if a.namespaced_name.namespace == namespace]
|
||||
@@ -787,9 +836,7 @@ class PlanLifecycleService:
|
||||
) -> list[Plan]:
|
||||
"""List plans with optional filtering.
|
||||
|
||||
Returns plans from the in-memory cache. When persistence is
|
||||
enabled the cache is populated during ``use_action`` and
|
||||
``get_plan`` calls.
|
||||
Merges plans from the in-memory cache and the persistence layer.
|
||||
|
||||
Args:
|
||||
namespace: Filter by namespace
|
||||
@@ -799,7 +846,33 @@ class PlanLifecycleService:
|
||||
Returns:
|
||||
List of matching Plans
|
||||
"""
|
||||
plans = list(self._plans.values())
|
||||
seen_ids: set[str] = set()
|
||||
plans: list[Plan] = []
|
||||
|
||||
for plan in self._plans.values():
|
||||
seen_ids.add(plan.identity.plan_id)
|
||||
plans.append(plan)
|
||||
|
||||
# Merge from persistence layer
|
||||
if self._persisted and self.unit_of_work is not None:
|
||||
try:
|
||||
with self.unit_of_work.transaction() as ctx:
|
||||
persisted = ctx.lifecycle_plans.list_plans(
|
||||
phase=phase.value if phase else None,
|
||||
namespace=namespace,
|
||||
project_name=project_name,
|
||||
)
|
||||
for plan in persisted:
|
||||
pid = plan.identity.plan_id
|
||||
if pid not in seen_ids:
|
||||
seen_ids.add(pid)
|
||||
plans.append(plan)
|
||||
self._plans[pid] = plan
|
||||
except Exception:
|
||||
self._logger.warning(
|
||||
"list_plans_persistence_fallback",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
if namespace:
|
||||
plans = [p for p in plans if p.namespaced_name.namespace == namespace]
|
||||
@@ -848,6 +921,12 @@ class PlanLifecycleService:
|
||||
)
|
||||
|
||||
# -- Pre-flight guardrail checks ----------------------------------
|
||||
# Ensure the plan's action is loaded into the in-memory cache
|
||||
if plan.action_name and plan.action_name not in self._actions:
|
||||
try:
|
||||
self.get_action(plan.action_name)
|
||||
except Exception:
|
||||
pass # Preflight will catch missing action below
|
||||
action_registry: dict[str, object] = {
|
||||
name: act for name, act in self._actions.items()
|
||||
}
|
||||
@@ -1077,6 +1156,388 @@ class PlanLifecycleService:
|
||||
|
||||
return plan
|
||||
|
||||
def run_execute(
|
||||
self,
|
||||
plan_id: str,
|
||||
project_path: Path | None = None,
|
||||
progress_callback: Any | None = None,
|
||||
) -> Plan:
|
||||
"""Run the full execute phase: start, invoke LLM, complete.
|
||||
|
||||
Orchestrates AI-powered code generation during the Execute phase:
|
||||
|
||||
1. Auto-advances through Strategize if the plan is still queued
|
||||
2. Transitions Execute from QUEUED → PROCESSING
|
||||
3. Resolves an AI provider via the actor system or registry
|
||||
4. Invokes the LLM via ``provider.generate_changes()``
|
||||
5. Writes generated files to a sandbox directory
|
||||
6. Records each file as a ``ChangeEntry`` in the changeset store
|
||||
7. Stores ``changeset_id`` and ``sandbox_refs`` on the plan
|
||||
8. Transitions Execute from PROCESSING → COMPLETE
|
||||
|
||||
Args:
|
||||
plan_id: The plan ULID
|
||||
project_path: Path to the project directory (for creating the
|
||||
legacy Project object). Defaults to current working directory.
|
||||
progress_callback: Optional callback for progress updates (0-100)
|
||||
|
||||
Returns:
|
||||
The updated Plan after execute completes
|
||||
"""
|
||||
from hashlib import sha256
|
||||
from pathlib import Path as _Path
|
||||
|
||||
from cleveragents.domain.models.core.change import (
|
||||
ChangeEntry,
|
||||
ChangeOperation,
|
||||
OperationType,
|
||||
)
|
||||
from cleveragents.domain.models.core.plan_legacy import Plan as LegacyPlan
|
||||
from cleveragents.domain.models.core.project_legacy import Project
|
||||
|
||||
plan = self.get_plan(plan_id)
|
||||
|
||||
# Auto-advance through strategize if needed
|
||||
if plan.phase == PlanPhase.STRATEGIZE:
|
||||
if plan.state == ProcessingState.QUEUED:
|
||||
self.start_strategize(plan_id)
|
||||
plan = self.get_plan(plan_id)
|
||||
if plan.state == ProcessingState.PROCESSING:
|
||||
self.complete_strategize(plan_id)
|
||||
plan = self.get_plan(plan_id)
|
||||
# complete_strategize may auto-progress to Execute via auto_progress()
|
||||
if plan.phase == PlanPhase.STRATEGIZE and plan.state == ProcessingState.COMPLETE:
|
||||
self.execute_plan(plan_id)
|
||||
plan = self.get_plan(plan_id)
|
||||
|
||||
# Start execute (QUEUED -> PROCESSING)
|
||||
if plan.phase == PlanPhase.EXECUTE and plan.state == ProcessingState.QUEUED:
|
||||
self.start_execute(plan_id)
|
||||
plan = self.get_plan(plan_id)
|
||||
|
||||
if plan.phase != PlanPhase.EXECUTE or plan.state != ProcessingState.PROCESSING:
|
||||
raise PlanError(
|
||||
f"Plan {plan_id} is not in Execute/PROCESSING state "
|
||||
f"(current: {plan.phase.value}/{plan.state})"
|
||||
)
|
||||
|
||||
# Resolve AI provider via actor system
|
||||
provider = self._resolve_provider(plan)
|
||||
if provider is None:
|
||||
self.fail_execute(plan_id, "No AI provider available")
|
||||
return self.get_plan(plan_id)
|
||||
|
||||
# Create legacy Plan adapter for the AI provider interface
|
||||
resolved_path = _Path(project_path) if project_path else _Path.cwd()
|
||||
legacy_project = Project(
|
||||
id=1,
|
||||
name=plan.namespaced_name.name,
|
||||
path=resolved_path,
|
||||
)
|
||||
# Include definition_of_done in the prompt so the LLM sees acceptance criteria
|
||||
prompt_parts = [plan.description]
|
||||
if plan.definition_of_done:
|
||||
prompt_parts.append(f"\nDefinition of Done:\n{plan.definition_of_done}")
|
||||
legacy_plan = LegacyPlan(
|
||||
id=1,
|
||||
project_id=1,
|
||||
name=plan.namespaced_name.name,
|
||||
prompt="\n".join(prompt_parts),
|
||||
)
|
||||
|
||||
self._logger.info(
|
||||
"Invoking AI provider for plan execution",
|
||||
plan_id=plan_id,
|
||||
provider=getattr(provider, "name", "unknown"),
|
||||
model=getattr(provider, "model_id", "unknown"),
|
||||
)
|
||||
|
||||
try:
|
||||
response = provider.generate_changes(
|
||||
project=legacy_project,
|
||||
plan=legacy_plan,
|
||||
contexts=[],
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
except Exception as exc:
|
||||
self._logger.error(
|
||||
"AI provider invocation failed",
|
||||
plan_id=plan_id,
|
||||
error=str(exc),
|
||||
)
|
||||
self.fail_execute(plan_id, f"AI provider error: {exc}")
|
||||
return self.get_plan(plan_id)
|
||||
|
||||
if response.error_message:
|
||||
self.fail_execute(plan_id, response.error_message)
|
||||
return self.get_plan(plan_id)
|
||||
|
||||
changes = response.changes or []
|
||||
|
||||
# Write generated files to sandbox directory
|
||||
sandbox_root = _Path(".cleveragents") / "sandbox" / plan_id
|
||||
sandbox_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Start a changeset for auditing
|
||||
changeset_id: str | None = None
|
||||
store = self._changeset_store
|
||||
if store is not None:
|
||||
try:
|
||||
changeset_id = store.start(plan_id)
|
||||
except Exception:
|
||||
self._logger.warning(
|
||||
"Failed to start changeset in store, continuing without audit",
|
||||
plan_id=plan_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
file_count = 0
|
||||
for change in changes:
|
||||
if change.operation not in (OperationType.CREATE, OperationType.MODIFY):
|
||||
continue
|
||||
content = change.new_content or ""
|
||||
file_path = change.file_path
|
||||
|
||||
# Write to sandbox
|
||||
sandbox_file = sandbox_root / file_path
|
||||
sandbox_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
sandbox_file.write_text(content)
|
||||
file_count += 1
|
||||
|
||||
# Record ChangeEntry in changeset store for audit trail
|
||||
if changeset_id and store is not None:
|
||||
try:
|
||||
from ulid import ULID
|
||||
|
||||
entry = ChangeEntry(
|
||||
plan_id=plan_id,
|
||||
resource_id=str(ULID()),
|
||||
tool_name="llm/generate",
|
||||
operation=ChangeOperation.CREATE,
|
||||
path=file_path,
|
||||
after_hash=sha256(content.encode()).hexdigest(),
|
||||
)
|
||||
store.record(changeset_id, entry)
|
||||
except Exception:
|
||||
self._logger.debug(
|
||||
"Failed to record changeset entry",
|
||||
path=file_path,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
# Store changeset_id and sandbox path on the plan
|
||||
plan = self.get_plan(plan_id)
|
||||
if changeset_id:
|
||||
plan.changeset_id = changeset_id
|
||||
plan.sandbox_refs = [str(sandbox_root.resolve())]
|
||||
plan.timestamps.updated_at = datetime.now()
|
||||
self._commit_plan(plan)
|
||||
|
||||
self._logger.info(
|
||||
"AI provider generated changes — files written to sandbox",
|
||||
plan_id=plan_id,
|
||||
sandbox=str(sandbox_root),
|
||||
file_count=file_count,
|
||||
changeset_id=changeset_id,
|
||||
model_used=response.model_used,
|
||||
token_count=response.token_count,
|
||||
)
|
||||
|
||||
# Complete execute (PROCESSING -> COMPLETE)
|
||||
return self.complete_execute(plan_id)
|
||||
|
||||
def run_apply(
|
||||
self,
|
||||
plan_id: str,
|
||||
project_path: Path | None = None,
|
||||
) -> tuple[Plan, int]:
|
||||
"""Run the full apply phase: copy sandbox files to the project.
|
||||
|
||||
Reads generated files from the sandbox directory (populated by
|
||||
``run_execute``) and writes them to the project root:
|
||||
|
||||
1. Transitions to Apply phase if in Execute/COMPLETE
|
||||
2. Reads sandbox path from ``plan.sandbox_refs``
|
||||
3. Copies each file from sandbox to project root
|
||||
4. Records apply completion
|
||||
5. Cleans up the sandbox directory
|
||||
|
||||
Args:
|
||||
plan_id: The plan ULID
|
||||
project_path: Path to write files in. Defaults to cwd.
|
||||
|
||||
Returns:
|
||||
Tuple of (updated Plan, number of files written)
|
||||
"""
|
||||
import shutil
|
||||
from pathlib import Path as _Path
|
||||
|
||||
plan = self.get_plan(plan_id)
|
||||
|
||||
# Transition to Apply if in Execute/COMPLETE
|
||||
if plan.phase == PlanPhase.EXECUTE and plan.state == ProcessingState.COMPLETE:
|
||||
self.apply_plan(plan_id)
|
||||
plan = self.get_plan(plan_id)
|
||||
|
||||
# Start apply if queued
|
||||
if plan.phase == PlanPhase.APPLY and plan.state == ProcessingState.QUEUED:
|
||||
self.start_apply(plan_id)
|
||||
plan = self.get_plan(plan_id)
|
||||
|
||||
if plan.phase != PlanPhase.APPLY or plan.state != ProcessingState.PROCESSING:
|
||||
raise PlanError(
|
||||
f"Plan {plan_id} is not in Apply/PROCESSING state "
|
||||
f"(current: {plan.phase.value}/{plan.state})"
|
||||
)
|
||||
|
||||
# Resolve sandbox path from plan
|
||||
sandbox_root: _Path | None = None
|
||||
if plan.sandbox_refs:
|
||||
candidate = _Path(plan.sandbox_refs[0])
|
||||
if candidate.is_dir():
|
||||
sandbox_root = candidate
|
||||
|
||||
if sandbox_root is None:
|
||||
self._logger.warning("No sandbox directory found", plan_id=plan_id)
|
||||
self.complete_apply(plan_id)
|
||||
return self.get_plan(plan_id), 0
|
||||
|
||||
project_root = (_Path(project_path) if project_path else _Path.cwd()).resolve()
|
||||
applied_count = 0
|
||||
|
||||
# Walk the sandbox and copy each file to the project root
|
||||
for sandbox_file in sandbox_root.rglob("*"):
|
||||
if not sandbox_file.is_file():
|
||||
continue
|
||||
try:
|
||||
relative = sandbox_file.relative_to(sandbox_root)
|
||||
target = (project_root / relative).resolve()
|
||||
|
||||
# Reject path traversal
|
||||
if not target.is_relative_to(project_root):
|
||||
self._logger.warning(
|
||||
"Path traversal rejected",
|
||||
path=str(relative),
|
||||
plan_id=plan_id,
|
||||
)
|
||||
continue
|
||||
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(sandbox_file.read_text())
|
||||
applied_count += 1
|
||||
except Exception as exc:
|
||||
self._logger.error(
|
||||
"Failed to apply file",
|
||||
file=str(sandbox_file),
|
||||
error=str(exc),
|
||||
plan_id=plan_id,
|
||||
)
|
||||
|
||||
self._logger.info(
|
||||
"Changes applied to filesystem from sandbox",
|
||||
plan_id=plan_id,
|
||||
sandbox=str(sandbox_root),
|
||||
applied_count=applied_count,
|
||||
)
|
||||
|
||||
# Complete apply
|
||||
self.complete_apply(plan_id)
|
||||
|
||||
# Clean up sandbox
|
||||
try:
|
||||
shutil.rmtree(sandbox_root)
|
||||
except Exception:
|
||||
self._logger.debug(
|
||||
"Failed to clean up sandbox",
|
||||
sandbox=str(sandbox_root),
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
# Clean up changeset store entries
|
||||
if self._changeset_store is not None and hasattr(self._changeset_store, "delete_for_plan"):
|
||||
try:
|
||||
self._changeset_store.delete_for_plan(plan_id)
|
||||
except Exception:
|
||||
self._logger.debug(
|
||||
"Failed to clean up changeset entries",
|
||||
plan_id=plan_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return self.get_plan(plan_id), applied_count
|
||||
|
||||
def _resolve_provider(self, plan: Plan) -> AIProviderInterface | None:
|
||||
"""Resolve an AI provider for executing the plan.
|
||||
|
||||
Resolution priority:
|
||||
1. Injected ``ai_provider`` (from container)
|
||||
2. Actor lookup via ``ActorService`` — extracts pre-stored
|
||||
``provider`` and ``model`` fields from the Actor domain object
|
||||
3. String-split fallback on ``plan.execution_actor``
|
||||
(e.g., ``"anthropic/claude-sonnet-4-20250514"``)
|
||||
4. Registry with default provider settings
|
||||
"""
|
||||
if self._ai_provider is not None:
|
||||
return self._ai_provider
|
||||
|
||||
# Try to get a registry
|
||||
registry = self._provider_registry
|
||||
if registry is None:
|
||||
try:
|
||||
from cleveragents.providers.registry import get_provider_registry
|
||||
|
||||
registry = get_provider_registry(self.settings)
|
||||
except Exception as exc:
|
||||
self._logger.error(
|
||||
"Failed to create provider registry", error=str(exc)
|
||||
)
|
||||
return None
|
||||
|
||||
actor_ref = plan.execution_actor or ""
|
||||
provider_type: str | None = None
|
||||
model_id: str | None = None
|
||||
|
||||
# Try actor-based resolution via ActorService
|
||||
if self._actor_service is not None and actor_ref:
|
||||
try:
|
||||
actor = self._actor_service.get_actor(actor_ref)
|
||||
provider_type = actor.provider or None
|
||||
model_id = actor.model or None
|
||||
self._logger.info(
|
||||
"Resolved provider via ActorService",
|
||||
actor=actor_ref,
|
||||
provider=provider_type,
|
||||
model=model_id,
|
||||
)
|
||||
except Exception:
|
||||
self._logger.debug(
|
||||
"Actor not found in database, falling back to string split",
|
||||
actor_ref=actor_ref,
|
||||
)
|
||||
|
||||
# Fallback: split the actor reference string
|
||||
if not provider_type:
|
||||
if "/" in actor_ref:
|
||||
provider_type, model_id = actor_ref.split("/", 1)
|
||||
else:
|
||||
provider_type = getattr(self.settings, "default_provider", "anthropic")
|
||||
model_id = getattr(self.settings, "default_model", None)
|
||||
|
||||
try:
|
||||
provider = registry.create_ai_provider(
|
||||
provider_type=provider_type,
|
||||
model_id=model_id or None,
|
||||
)
|
||||
return provider
|
||||
except Exception as exc:
|
||||
self._logger.error(
|
||||
"Failed to create AI provider",
|
||||
provider=provider_type,
|
||||
model=model_id,
|
||||
error=str(exc),
|
||||
)
|
||||
return None
|
||||
|
||||
def apply_plan(self, plan_id: str) -> Plan:
|
||||
"""Transition plan from Execute to Apply phase.
|
||||
|
||||
|
||||
@@ -39,6 +39,13 @@ if TYPE_CHECKING:
|
||||
|
||||
_logger = structlog.get_logger(__name__)
|
||||
|
||||
# SQLAlchemy session type for CommittingSessionService
|
||||
try:
|
||||
from sqlalchemy.orm import Session as SASession, sessionmaker as SASessionMaker
|
||||
except ImportError: # pragma: no cover
|
||||
SASession = Any # type: ignore[assignment,misc]
|
||||
SASessionMaker = Any # type: ignore[assignment,misc]
|
||||
|
||||
|
||||
class PersistentSessionService(SessionService):
|
||||
"""Database-backed session service.
|
||||
@@ -338,3 +345,66 @@ class PersistentSessionService(SessionService):
|
||||
session.updated_at = datetime.now()
|
||||
|
||||
self._session_repo.update(session)
|
||||
|
||||
|
||||
class CommittingSessionService(PersistentSessionService):
|
||||
"""Session service that commits after each mutating operation.
|
||||
|
||||
For CLI usage where there is no UoW wrapper managing transactions.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_repo: SessionRepository,
|
||||
message_repo: SessionMessageRepository,
|
||||
session_factory: Any,
|
||||
) -> None:
|
||||
super().__init__(session_repo, message_repo)
|
||||
self._sa_factory = session_factory
|
||||
|
||||
def _commit(self) -> None:
|
||||
"""Commit the current database session."""
|
||||
try:
|
||||
sa_session = self._sa_factory()
|
||||
sa_session.commit()
|
||||
except Exception:
|
||||
import structlog
|
||||
|
||||
structlog.get_logger(__name__).warning(
|
||||
"session commit failed", exc_info=True
|
||||
)
|
||||
|
||||
def create(self, actor_name: str | None = None) -> Session:
|
||||
result = super().create(actor_name)
|
||||
self._commit()
|
||||
return result
|
||||
|
||||
def delete(self, session_id: str) -> None:
|
||||
super().delete(session_id)
|
||||
self._commit()
|
||||
|
||||
def append_message(
|
||||
self,
|
||||
session_id: str,
|
||||
role: MessageRole,
|
||||
content: str,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> SessionMessage:
|
||||
result = super().append_message(session_id, role, content, metadata)
|
||||
self._commit()
|
||||
return result
|
||||
|
||||
def import_session(self, data: dict[str, Any]) -> Session:
|
||||
result = super().import_session(data)
|
||||
self._commit()
|
||||
return result
|
||||
|
||||
def update_token_usage(
|
||||
self,
|
||||
session_id: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cost: float,
|
||||
) -> None:
|
||||
super().update_token_usage(session_id, input_tokens, output_tokens, cost)
|
||||
self._commit()
|
||||
|
||||
@@ -1530,30 +1530,34 @@ def execute_plan(
|
||||
try:
|
||||
service = _get_lifecycle_service()
|
||||
|
||||
if not plan_id:
|
||||
# Try to find the only plan in strategize phase
|
||||
from cleveragents.domain.models.core.plan import (
|
||||
PlanPhase,
|
||||
ProcessingState,
|
||||
)
|
||||
from cleveragents.domain.models.core.plan import (
|
||||
PlanPhase,
|
||||
ProcessingState,
|
||||
)
|
||||
|
||||
if not plan_id:
|
||||
# Try to find plans in strategize phase (complete or queued)
|
||||
plans = service.list_plans(phase=PlanPhase.STRATEGIZE)
|
||||
complete_plans = [p for p in plans if p.state == ProcessingState.COMPLETE]
|
||||
if len(complete_plans) == 0:
|
||||
eligible_plans = [
|
||||
p
|
||||
for p in plans
|
||||
if p.state in (ProcessingState.COMPLETE, ProcessingState.QUEUED)
|
||||
]
|
||||
if len(eligible_plans) == 0:
|
||||
console.print(
|
||||
"[yellow]No plans ready for execution.[/yellow]\n"
|
||||
"Plans must be in Strategize phase with 'complete' state."
|
||||
"Plans must be in Strategize phase."
|
||||
)
|
||||
raise typer.Abort()
|
||||
if len(complete_plans) > 1:
|
||||
if len(eligible_plans) > 1:
|
||||
console.print(
|
||||
"[yellow]Multiple plans ready for execution. "
|
||||
"Please specify a plan ID.[/yellow]"
|
||||
)
|
||||
for p in complete_plans:
|
||||
for p in eligible_plans:
|
||||
console.print(f" • {p.identity.plan_id}: {p.namespaced_name}")
|
||||
raise typer.Abort()
|
||||
plan_id = complete_plans[0].identity.plan_id
|
||||
plan_id = eligible_plans[0].identity.plan_id
|
||||
|
||||
# Apply execution environment override if provided
|
||||
if execution_environment:
|
||||
@@ -1580,18 +1584,24 @@ def execute_plan(
|
||||
)
|
||||
raise typer.Abort()
|
||||
|
||||
# Execute the plan
|
||||
plan = service.execute_plan(plan_id)
|
||||
# Run the full execute phase with real LLM invocation
|
||||
console.print("[dim]Running execute phase with AI provider...[/dim]")
|
||||
plan = service.run_execute(plan_id)
|
||||
|
||||
if fmt != OutputFormat.RICH.value:
|
||||
data = _plan_spec_dict(plan)
|
||||
console.print(format_output(data, fmt))
|
||||
else:
|
||||
_print_lifecycle_plan(plan, title="Plan Executing")
|
||||
console.print(
|
||||
"\n[dim]Plan is now in Execute phase (queued). "
|
||||
"Run 'agents plan apply <id>' when execution is complete.[/dim]"
|
||||
)
|
||||
_print_lifecycle_plan(plan, title="Plan Executed")
|
||||
if plan.state == ProcessingState.ERRORED:
|
||||
console.print(
|
||||
f"\n[red]Execute failed:[/red] {plan.error_message}"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
"\n[dim]Execute phase complete. "
|
||||
"Run 'agents plan lifecycle-apply <id>' to write files.[/dim]"
|
||||
)
|
||||
|
||||
except InvalidPhaseTransitionError as e:
|
||||
console.print(f"[red]Invalid transition:[/red] {e}")
|
||||
@@ -1610,6 +1620,14 @@ def lifecycle_apply_plan(
|
||||
str | None,
|
||||
typer.Argument(help="Plan ID to apply (optional if only one plan)"),
|
||||
] = None,
|
||||
output_dir: Annotated[
|
||||
str | None,
|
||||
typer.Option(
|
||||
"--output-dir",
|
||||
"-o",
|
||||
help="Directory to write generated files into (default: current directory).",
|
||||
),
|
||||
] = None,
|
||||
fmt: Annotated[
|
||||
str,
|
||||
typer.Option(
|
||||
@@ -1632,30 +1650,34 @@ def lifecycle_apply_plan(
|
||||
try:
|
||||
service = _get_lifecycle_service()
|
||||
|
||||
if not plan_id:
|
||||
# Try to find the only plan in execute phase
|
||||
from cleveragents.domain.models.core.plan import (
|
||||
PlanPhase,
|
||||
ProcessingState,
|
||||
)
|
||||
from cleveragents.domain.models.core.plan import (
|
||||
PlanPhase,
|
||||
ProcessingState,
|
||||
)
|
||||
|
||||
if not plan_id:
|
||||
# Try to find plans in execute phase (complete or queued)
|
||||
plans = service.list_plans(phase=PlanPhase.EXECUTE)
|
||||
complete_plans = [p for p in plans if p.state == ProcessingState.COMPLETE]
|
||||
if len(complete_plans) == 0:
|
||||
eligible_plans = [
|
||||
p
|
||||
for p in plans
|
||||
if p.state in (ProcessingState.COMPLETE, ProcessingState.QUEUED)
|
||||
]
|
||||
if len(eligible_plans) == 0:
|
||||
console.print(
|
||||
"[yellow]No plans ready for apply.[/yellow]\n"
|
||||
"Plans must be in Execute phase with 'complete' state."
|
||||
"Plans must be in Execute phase."
|
||||
)
|
||||
raise typer.Abort()
|
||||
if len(complete_plans) > 1:
|
||||
if len(eligible_plans) > 1:
|
||||
console.print(
|
||||
"[yellow]Multiple plans ready for apply. "
|
||||
"Please specify a plan ID.[/yellow]"
|
||||
)
|
||||
for p in complete_plans:
|
||||
for p in eligible_plans:
|
||||
console.print(f" • {p.identity.plan_id}: {p.namespaced_name}")
|
||||
raise typer.Abort()
|
||||
plan_id = complete_plans[0].identity.plan_id
|
||||
plan_id = eligible_plans[0].identity.plan_id
|
||||
|
||||
# Fail-fast: read-only plans must not enter Apply phase
|
||||
pre_plan = service.get_plan(plan_id)
|
||||
@@ -1665,17 +1687,27 @@ def lifecycle_apply_plan(
|
||||
)
|
||||
raise typer.Abort()
|
||||
|
||||
# Apply the plan
|
||||
plan = service.apply_plan(plan_id)
|
||||
# Run the full apply phase — writes files to disk
|
||||
from pathlib import Path
|
||||
|
||||
target_path: Path | None = None
|
||||
if output_dir:
|
||||
target_path = Path(output_dir).resolve()
|
||||
target_path.mkdir(parents=True, exist_ok=True)
|
||||
console.print(
|
||||
f"[dim]Applying changes to {target_path} ...[/dim]"
|
||||
)
|
||||
else:
|
||||
console.print("[dim]Applying changes to filesystem...[/dim]")
|
||||
plan, applied_count = service.run_apply(plan_id, project_path=target_path)
|
||||
|
||||
if fmt != OutputFormat.RICH.value:
|
||||
data = _plan_spec_dict(plan)
|
||||
console.print(format_output(data, fmt))
|
||||
else:
|
||||
_print_lifecycle_plan(plan, title="Plan Applying")
|
||||
_print_lifecycle_plan(plan, title="Plan Applied")
|
||||
console.print(
|
||||
"\n[dim]Plan is now in Apply phase (queued). "
|
||||
"Changes will be applied to the project(s).[/dim]"
|
||||
f"\n[green]{applied_count} file(s) written to disk.[/green]"
|
||||
)
|
||||
|
||||
except InvalidPhaseTransitionError as e:
|
||||
|
||||
@@ -14,7 +14,7 @@ import logging
|
||||
import sys
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any, cast
|
||||
from typing import Annotated, Any
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
@@ -57,11 +57,31 @@ def _get_session_service() -> SessionService:
|
||||
if _service is not None:
|
||||
return _service
|
||||
|
||||
from cleveragents.application.container import get_container
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
container = get_container()
|
||||
svc = cast(SessionService, container.session_service())
|
||||
_service = svc
|
||||
from cleveragents.application.container import get_database_url
|
||||
from cleveragents.application.services.session_service import (
|
||||
CommittingSessionService,
|
||||
)
|
||||
from cleveragents.infrastructure.database.repositories import (
|
||||
SessionMessageRepository,
|
||||
SessionRepository,
|
||||
)
|
||||
|
||||
database_url = get_database_url()
|
||||
engine = create_engine(database_url, echo=False)
|
||||
raw_factory = sessionmaker(bind=engine, expire_on_commit=False)
|
||||
|
||||
# Repos need to share a single session so flush + commit work together.
|
||||
_shared_session = raw_factory()
|
||||
|
||||
def shared_factory():
|
||||
return _shared_session
|
||||
|
||||
session_repo = SessionRepository(session_factory=shared_factory)
|
||||
message_repo = SessionMessageRepository(session_factory=shared_factory)
|
||||
_service = CommittingSessionService(session_repo, message_repo, shared_factory)
|
||||
return _service
|
||||
|
||||
|
||||
|
||||
@@ -601,6 +601,7 @@ def main(args: list[str] | None = None) -> int:
|
||||
"resource", # Resource registry management
|
||||
"skill", # Skill management
|
||||
"lsp", # LSP server management
|
||||
"audit", # Audit log management
|
||||
"cleanup", # Garbage collection and cleanup
|
||||
"config", # Configuration management
|
||||
"session", # Session management
|
||||
|
||||
@@ -10,6 +10,15 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
from cleveragents.domain.models.core.retry_policy import RetryStrategy
|
||||
|
||||
# Load .env into os.environ so provider API keys (ANTHROPIC_API_KEY, etc.)
|
||||
# are available to _apply_external_env_overrides and LangChain constructors.
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv(override=False)
|
||||
except ImportError: # pragma: no cover - python-dotenv is optional
|
||||
pass
|
||||
|
||||
# pydantic v2: use string for extra mode
|
||||
EXTRA_IGNORE = "ignore"
|
||||
STRUCTLOG_ENABLED = False
|
||||
@@ -27,7 +36,11 @@ class Settings(BaseSettings):
|
||||
"""Application runtime configuration backed by environment variables."""
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="cleveragents_", case_sensitive=False, extra="ignore"
|
||||
env_prefix="cleveragents_",
|
||||
case_sensitive=False,
|
||||
extra="ignore",
|
||||
env_file=".env",
|
||||
env_file_encoding="utf-8",
|
||||
)
|
||||
|
||||
_instance: ClassVar[Settings | None] = None
|
||||
|
||||
@@ -158,22 +158,31 @@ class MigrationRunner:
|
||||
def get_pending_migrations(self) -> list[str]:
|
||||
"""Get list of pending migrations.
|
||||
|
||||
Uses Alembic's built-in head-comparison to correctly determine
|
||||
which revisions have not yet been applied, handling branches and
|
||||
merges properly.
|
||||
|
||||
Returns:
|
||||
List of migration revision IDs that haven't been applied yet
|
||||
"""
|
||||
# Get current revision
|
||||
current_rev = self.get_current_revision()
|
||||
|
||||
# Get all revisions from Alembic scripts
|
||||
script_dir = ScriptDirectory.from_config(self.alembic_cfg)
|
||||
head = script_dir.get_current_head()
|
||||
|
||||
# Get revisions that come after the current revision
|
||||
# If current revision matches head, nothing is pending
|
||||
if current_rev == head:
|
||||
return []
|
||||
|
||||
# If no migrations applied yet, everything is pending
|
||||
if current_rev is None:
|
||||
return [rev.revision for rev in reversed(list(script_dir.walk_revisions()))]
|
||||
|
||||
# Walk from head backwards, collecting revisions until we hit current
|
||||
pending: list[str] = []
|
||||
for rev in script_dir.walk_revisions():
|
||||
if current_rev is None or rev.revision != current_rev:
|
||||
pending.append(rev.revision)
|
||||
if rev.revision == current_rev:
|
||||
break
|
||||
if rev.revision == current_rev:
|
||||
break
|
||||
pending.append(rev.revision)
|
||||
|
||||
# Reverse to get chronological order
|
||||
return list(reversed(pending))
|
||||
|
||||
@@ -79,10 +79,18 @@ class LangChainChatProvider(AIProviderInterface):
|
||||
def name(self) -> str: # pragma: no cover - simple accessor
|
||||
return self._name
|
||||
|
||||
@name.setter
|
||||
def name(self, value: str) -> None:
|
||||
self._name = value
|
||||
|
||||
@property
|
||||
def model_id(self) -> str: # pragma: no cover - simple accessor
|
||||
return self._model_id
|
||||
|
||||
@model_id.setter
|
||||
def model_id(self, value: str) -> None:
|
||||
self._model_id = value
|
||||
|
||||
def generate_changes(
|
||||
self,
|
||||
project: Project,
|
||||
|
||||
@@ -650,6 +650,12 @@ class ProviderRegistry:
|
||||
# Create factory function for the LLM
|
||||
def llm_factory(mid: str) -> BaseLanguageModel:
|
||||
factory_kwargs: dict[str, Any] = {"max_retries": max_retries}
|
||||
# Pass API key from settings so .env keys work without shell export
|
||||
key_attr = self.PROVIDER_KEY_ATTRS.get(provider_type)
|
||||
if key_attr:
|
||||
api_key = getattr(self._settings, key_attr, None)
|
||||
if api_key:
|
||||
factory_kwargs["api_key"] = api_key
|
||||
return self._create_provider_llm(
|
||||
provider_type,
|
||||
mid,
|
||||
|
||||
Reference in New Issue
Block a user