Compare commits

...

25 Commits

Author SHA1 Message Date
aditya 6158046eae test: sim12, sim14, sim17 — verify multi-file DoD compliance after prompt fix
CI / lint (pull_request) Failing after 17s
CI / benchmark-publish (pull_request) Has been skipped
CI / build (pull_request) Successful in 15s
CI / quality (pull_request) Successful in 40s
CI / security (pull_request) Failing after 40s
CI / typecheck (pull_request) Failing after 1m0s
CI / coverage (pull_request) Has been skipped
CI / benchmark-regression (pull_request) Has been skipped
CI / e2e_tests (pull_request) Successful in 56s
CI / integration_tests (pull_request) Failing after 3m35s
CI / unit_tests (pull_request) Failing after 3m48s
CI / docker (pull_request) Has been skipped
All 3 simulations tested with the fixed generate prompt that passes
the original DoD to the LLM. Results:

- SIM12 (Rate Limiter): 7 files — algorithms.py, store.py, middleware.py,
  config.py, example_app.py, requirements.txt, __init__.py
- SIM14 (JSON Validator): 6 files — validator.py, types.py, errors.py,
  cli.py, requirements.txt, __init__.py
- SIM17 (API Tester): 7 files — runner.py, assertions.py, variables.py,
  reporter.py, cli.py, requirements.txt, test_suite_example.yaml

All Python files pass compile() syntax check. All DoD files present.
Files written directly via --output-dir, no manual cleanup needed.
2026-03-13 16:46:27 +00:00
aditya ce948da5ff fix: pass original prompt with DoD to code generation LLM
The generate prompt only received the analyze_requirements summary,
which often compressed the file-by-file DoD into a brief paragraph.
The LLM then produced 1 file instead of the 5-6 specified.

Fix: added {original_prompt} variable to the generate prompt template
so the LLM sees both the architecture analysis AND the original DoD
with explicit file listings. Also strengthened the instruction to
"output ALL files mentioned in the Definition of Done."

Verified: sim10 (CSV Analyzer) now produces 6 files matching DoD
(loader.py, analyzer.py, reporter.py, cli.py, requirements.txt).
2026-03-13 16:46:26 +00:00
aditya fea2cc2649 feat: add --output-dir flag to lifecycle-apply command
Files were always written to CWD, requiring manual mv and cleanup
(README.md overwrites, stray files in project root). Now users can
specify where generated files should go:

  python -m cleveragents plan lifecycle-apply PLAN_ID --output-dir simulations/sim9

The directory is created automatically if it doesn't exist.
Updated sim_test_commands.md to use the new flag.
2026-03-13 16:46:26 +00:00
aditya 7d646d3439 test: sim9 multi-file RAG application with src/ subdirectory (10 files)
First simulation testing subdirectory output support. RAG app with
separate modules: document_loader, chunker, embeddings, vector_store,
retriever, generator — plus config.py and main.py CLI entry point.
All Python files pass syntax check. 9/10 DoD checks pass.

Also adds sim_test_commands.md with 9 new simulation definitions
(SIM9-SIM17) for manual terminal testing.
2026-03-13 16:46:25 +00:00
aditya 14d390a660 test: sim8 basic RAG application — end-to-end V3 simulation (2 files)
User-triggered simulation via terminal commands. RAG app with OpenAI
embeddings, FAISS vector store, document chunking, and CLI interface.
2026-03-13 16:46:25 +00:00
aditya 07f741fafa fix: broken pending migration detection causing repeated prompts
get_pending_migrations() had a logic bug where the break condition
(rev.revision == current_rev) was inside the (rev.revision != current_rev)
guard, making it unreachable. This caused every CLI command to falsely
detect pending migrations and prompt for approval even when the DB was
already at head.
2026-03-13 16:46:24 +00:00
aditya 07d0bb9da6 test: re-run all 7 simulations with improved pipeline
All sims re-run after pipeline fixes (better prompts, real syntax
validation, fenced-block extraction fix). Results: 7/7 PASS,
20 files total, all Python files pass compile() syntax check.

Includes run_all_sims.py runner script and rag-basic action config.
2026-03-13 16:46:24 +00:00
aditya 59ea0d2092 docs: annotate comparative review with implementation responses
Add [RESPONSE] annotations to each finding in the comparative review:
- Findings A, B, C, E: acknowledged as valid with action items
- Finding D (binary apply): noted as low priority / theoretical
- Finding F (simulation artifacts): corrected — reviewer checked raw
  LLM outputs (sim*-<name>/generated.py) instead of extracted files
  (sim*/todo.py etc.) which all pass py_compile

Also fix CommittingSessionService._commit() to log exceptions instead
of silently swallowing them (valid Finding E).
2026-03-13 16:46:23 +00:00
aditya a2c859aef6 fix: improve code generation pipeline prompts, validation, and file extraction
Problems fixed:
- Generation prompt was vague ("produce production-ready code") with no
  output format specification. LLM had to guess the multi-file format.
- Validation was fake: `is_valid = "PASS" in text or len(code) > 10`
  auto-passed any code longer than 10 chars.
- Single-file fallback used keyword guessing ("test" -> test_generated.py)
  instead of parsing the filename from LLM output.

Changes:
- Rewrite analyze_prompt with structured sections (Files to Create,
  Dependencies, Key Requirements, Architecture Notes)
- Rewrite generate_prompt with explicit multi-file output format rules
  (**filename.ext** + fenced code block) and quality constraints
- Rewrite validate_prompt with specific review checklist
- Replace fake validation with real Python syntax check (compile()) +
  LLM review where first line must be PASS or FAIL
- Add _extract_single_file() for robust single-file fallback that
  parses the filename from the LLM output header
- Remove keyword-based path guessing (test/error/generated.py)
- Re-ran SIM1: now produces todo.py (matching DoD) instead of main.py
2026-03-13 16:46:23 +00:00
aditya ac1e17b8bd test: sim7 URL shortener — end-to-end V3 simulation (4 files)
Flask URL shortener with base62 encoding, POST /shorten, GET /<code>
redirect, input validation, and HTML templates (index + error page).
Multi-file extraction correctly handled nested templates/ directory.
2026-03-13 16:46:22 +00:00
aditya 527fccdf6b test: sim6 weather CLI dashboard — end-to-end V3 simulation (2 files)
New simulation: OpenWeatherMap CLI with WeatherData dataclass, argparse,
formatted terminal output, and retry-enabled HTTP client.
Multi-file extraction correctly produced weather.py + requirements.txt.
2026-03-13 16:46:22 +00:00
aditya 9e89c031fd test: sim2 bookstore REST API extracted output (4 files) 2026-03-13 16:46:22 +00:00
aditya 82078bd9e8 chore: add docs/camodif.md to .gitignore 2026-03-13 16:46:18 +00:00
aditya 2f4f4c9b49 fix: multi-file extraction skips headers inside fenced code blocks
_extract_files_from_markdown now builds a set of fenced-block character
ranges and filters out any FILE_HEADER_RE matches that fall within them.
This prevents dependency lines like "requests>=2.31.0" from being
misidentified as file headers, which was causing requirements.txt to be
missed in sim4's output. Re-extracted sim4 with correct scraper.py and
requirements.txt.
2026-03-13 16:45:05 +00:00
aditya ada913179d refactor: use existing infrastructure for v3 plan lifecycle
Replace custom JSON file persistence, string-split provider resolution,
and raw change tracking with the project's existing infrastructure:
- ActorService for provider/model resolution from DB
- ChangeSetStore (InMemoryChangeSetStore) for audit-grade change tracking
- Sandbox directories for cross-process file persistence
- Definition of Done included in LLM prompts
2026-03-13 16:45:04 +00:00
aditya ca33ce2f70 feat: v3 simulation action configs and improved multi-file extraction
Add 5 action YAML configs for v3 simulation stress tests (todo-cli,
bookstore-api, websocket-chat, hn-scraper, flask-auth). Update
multi-file markdown extraction regex to also match bold-only headers
(**filename.py**) without requiring # prefix.
2026-03-13 16:45:04 +00:00
aditya 2c13e747b4 feat: wire real LLM calls into v3 plan lifecycle
The v3 lifecycle (action create → plan use → plan execute → plan
lifecycle-apply) now invokes the actual AI provider during the execute
phase and writes generated files to disk during the apply phase.

Key changes:
- Add ai_provider and provider_registry to PlanLifecycleService
- Add run_execute() method that resolves provider, creates legacy Plan
  adapter, calls generate_changes(), and persists changeset to JSON
- Add run_apply() method that loads persisted changeset and writes
  files to filesystem with path traversal protection
- Wire ai_provider and provider_registry in DI container
- Update CLI execute/lifecycle-apply commands to use new methods
- Load .env via python-dotenv so API keys are available without shell
  export
- Pass API keys from settings into LLM factory kwargs
2026-03-13 16:45:03 +00:00
aditya a4918ef0e6 feat: multi-file code generation from LLM output
PlanGenerationGraph._generate_plan() now parses markdown-formatted LLM
output to extract individual files from ## headers and fenced code blocks.
Previously, all LLM output was crammed into a single generated.py file.

Now: ## main.py + ```python``` blocks are parsed into separate Change
objects, each written to its own file path. Supports nested paths like
templates/base.html, numbered headers (## 1. file.py), backtick-wrapped
and bold filenames. Falls back to single-file behavior when < 2 file
headers are detected.

Tested against 5 real Anthropic LLM outputs — correctly extracts 2-11
files per generation including .py, .txt, .html, .css, and .md files.
2026-03-13 16:45:03 +00:00
aditya e1e85ed9c7 test: 5 real-LLM simulation outputs from end-to-end workflow stress test
Each simulation ran the full legacy lifecycle (init → tell → build → apply)
with real Anthropic Claude API calls:

1. sim1-todo-cli: Click-based todo app (add/list/delete)
2. sim2-bookstore-api: FastAPI CRUD with SQLite (5 code files)
3. sim3-websocket-chat: async websocket chat server
4. sim4-hn-scraper: HN top stories scraper with JSON output
5. sim5-flask-auth: Flask + Flask-Login authentication app

All 5 generated correct, production-quality Python code matching
the requested task specifications.
2026-03-13 16:45:02 +00:00
aditya a1b65057a9 fix: add property setters to LangChainChatProvider for name and model_id
PlanService._resolve_ai_provider_for_actor() sets name and model_id on the
provider instance, but LangChainChatProvider only had read-only properties.
This caused AttributeError when building plans with a real LLM provider.

Also adds .gitignore (protecting .env with API keys) and example action YAML.
2026-03-13 16:44:59 +00:00
aditya d468e36e59 fix: session persistence uses shared SQLAlchemy session for proper commit
The session factory was creating new sessions on each call, so flush() in
the repo and commit() in CommittingSessionService operated on different
sessions. Now uses a shared session instance so changes are properly
committed to the database.
2026-03-13 16:42:00 +00:00
aditya a6e1151209 fix: add 'audit' to valid CLI commands list
The audit command was registered as a subcommand but missing from the
valid_cmds whitelist in main(), causing 'Invalid command' errors.
2026-03-13 16:39:15 +00:00
aditya 444cbbfcc6 fix: session CLI now persists to database via CommittingSessionService
The session CLI was using container.db() which didn't exist. Replaced with
proper SQLAlchemy session factory pattern. Created CommittingSessionService
that wraps PersistentSessionService with auto-commit after each mutating
operation, since the CLI doesn't use a UoW wrapper for transactions.
2026-03-13 16:39:08 +00:00
aditya 7cb5a542cb fix: v3 plan lifecycle auto-advances through phases and persists correctly
- execute command now auto-advances plans from strategize/queued through
  start_strategize and complete_strategize before executing
- lifecycle-apply command auto-advances from execute/queued through
  start_execute and complete_execute before applying
- list_plans now queries the persistence layer (same fix as list_actions)
- start_strategize loads the action from DB into cache before preflight
  checks, fixing "Action not found in registry" when action was only in DB

The full v3 lifecycle now works end-to-end:
  action create → plan use → plan execute → plan lifecycle-apply
2026-03-13 16:31:56 +00:00
aditya cdf13f842d fix: use container-wired PlanLifecycleService for persistence in CLI commands
Both action.py and plan.py CLI commands were creating PlanLifecycleService
with only settings (no UnitOfWork), causing in-memory-only storage that lost
data between CLI invocations. Now uses container.plan_lifecycle_service()
which injects a real UnitOfWork for database persistence.

Also fixed list_actions() to query the persistence layer instead of only
reading from the in-memory cache, which was always empty on fresh service
instances.
2026-03-13 16:31:46 +00:00
91 changed files with 9869 additions and 147 deletions
+9
View File
@@ -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
+7
View File
@@ -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.
+15
View File
@@ -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).
+11
View File
@@ -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.
+14
View File
@@ -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.
+16
View File
@@ -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).
+15
View File
@@ -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).
+16
View File
@@ -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).
+13
View File
@@ -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.
+12
View File
@@ -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.
+12
View File
@@ -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.
+12
View File
@@ -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.
+15
View File
@@ -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.
+15
View File
@@ -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.
+21
View File
@@ -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
View File
@@ -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}")
+316
View File
@@ -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.
+1
View File
@@ -0,0 +1 @@
[]
+253
View File
@@ -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()
+7
View File
@@ -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"
+246
View File
@@ -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()
}
+187
View File
@@ -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()
+192
View File
@@ -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
+264
View File
@@ -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}")
+4
View File
@@ -0,0 +1,4 @@
pandas>=1.3.0
numpy>=1.20.0
chardet>=4.0.0
scipy>=1.7.0
+19
View File
@@ -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"
]
+151
View File
@@ -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
+25
View File
@@ -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")
+105
View File
@@ -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)
+94
View File
@@ -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
+1
View File
@@ -0,0 +1 @@
flask>=2.0.0
+86
View File
@@ -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()
+38
View File
@@ -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"
]
+109
View File
@@ -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()
+35
View File
@@ -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
}
+2
View File
@@ -0,0 +1,2 @@
# No external dependencies required
# This JSON schema validator uses only Python standard library
+208
View File
@@ -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__}"
)
+152
View File
@@ -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
+185
View File
@@ -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}")
+98
View File
@@ -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()
+137
View File
@@ -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))
+3
View File
@@ -0,0 +1,3 @@
requests>=2.28.0
pyyaml>=6.0
jsonpath-ng>=1.5.0
+186
View File
@@ -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
+153
View File
@@ -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
+121
View File
@@ -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
+670
View File
@@ -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.
+120
View File
@@ -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)
+3
View File
@@ -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.
+423
View File
@@ -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>
+1
View File
@@ -0,0 +1 @@
websockets>=11.0.0
+139
View File
@@ -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")
+320
View File
@@ -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.
+3
View File
@@ -0,0 +1,3 @@
requests>=2.28.0
beautifulsoup4>=4.11.0
lxml>=4.9.0
+251
View File
@@ -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()
+552
View File
@@ -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`
+361
View File
@@ -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)
+119
View File
@@ -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
+12
View File
@@ -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))
+36
View File
@@ -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
+45
View File
@@ -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}>'
+5
View File
@@ -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
+1
View File
@@ -0,0 +1 @@
requests>=2.28.0
+193
View File
@@ -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())
+121
View File
@@ -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)
+2
View File
@@ -0,0 +1,2 @@
flask>=2.0.0
validators>=0.20.0
+161
View File
@@ -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>
+376
View File
@@ -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())
+5
View File
@@ -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
+6
View File
@@ -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
+73
View File
@@ -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
+228
View File
@@ -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()
+5
View File
@@ -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
View File
+72
View File
@@ -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
+41
View File
@@ -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
+68
View File
@@ -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([])
+130
View File
@@ -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)}"
+60
View File
@@ -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
+125
View File
@@ -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()
+247 -92
View File
@@ -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.",
},
}
+27
View File
@@ -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()
+68 -36
View File
@@ -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:
+25 -5
View File
@@ -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
+1
View File
@@ -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
+14 -1
View File
@@ -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,
+6
View File
@@ -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,