refactor: improve error handling and remove hardcoded credentials
- Changed bare except clauses to specific exception types in web3_utils.py, testing.py, messages.py, and message_storage.py - Replaced print() calls with logger in testing.py, agent_discovery.py, compliance_agent.py, coordinator.py, trading_agent.py, keys.py, escrow.py, persistent_spending_tracker.py, sync_cli.py, and client.py - Added logger initialization using get_logger(__name__) in compliance_agent.py, coordinator.py, trading_agent.py, keys.py, escrow.py, persistent_spending_tracker.py, and client.py - Removed hardcoded secret
This commit is contained in:
aitbc
97
.hermes/plans/2026-05-12_104500-coordinator-decomposition.md
Normal file
97
.hermes/plans/2026-05-12_104500-coordinator-decomposition.md
Normal file
@@ -0,0 +1,97 @@
|
||||
|
||||
# Coordinator-API Decomposition Plan
|
||||
|
||||
## Current State
|
||||
- **1 monolith**: apps/coordinator-api/src/app/
|
||||
- 89 service files, 46,594 LOC
|
||||
- 53 routers
|
||||
- 51 files over 500 LOC
|
||||
- Largest: agent_integration.py (1,159 LOC)
|
||||
|
||||
## Decomposition Strategy: Bounded Contexts
|
||||
|
||||
Based on domain analysis, split into 7 microservices:
|
||||
|
||||
1. **agent-management** (agent lifecycle, performance, communication)
|
||||
2. **blockchain** (chain operations, transactions, smart contracts)
|
||||
3. **computing** (GPU, resources, marketplace for compute)
|
||||
4. **enterprise** (integration, scalability, compliance)
|
||||
5. **identity** (authentication, authorization, agents identity)
|
||||
6. **payment** (billing, transactions, financial operations)
|
||||
7. **ai-models** (AI services, RL, multi-modal fusion)
|
||||
|
||||
Each will be a separate FastAPI app with:
|
||||
- Its own routers/, services/, models/
|
||||
- Shared libraries: app.core.config, app.core.logging, app.core.database
|
||||
- Independent systemd service
|
||||
- Clear API boundaries
|
||||
|
||||
## Implementation Phases
|
||||
|
||||
### Phase 1: Infrastructure Setup (Week 1-2)
|
||||
- Create apps/ directory structure: agent-management/, blockchain/, etc.
|
||||
- Create shared core library: apps/coordinator-api/src/app/core/
|
||||
- Extract common config, logging, DB session, exceptions
|
||||
- Update pyproject.toml to support multiple packages
|
||||
|
||||
### Phase 2: Extract Agent Management (Week 2-3)
|
||||
- Move agent_*.py, agent_service_marketplace.py -> agent-management
|
||||
- Move agent_communication.py, agent_performance_service.py -> agent-management
|
||||
- Create new systemd service for agent-management
|
||||
- Update reverse proxy (nginx) routes
|
||||
|
||||
### Phase 3: Extract Blockchain (Week 3-4)
|
||||
- Move blockchain_context.py, contract_service.py, transaction_service.py -> blockchain
|
||||
- Move escrow.py, persistent_spending_tracker.py, etc.
|
||||
- Create blockchain systemd service
|
||||
|
||||
### Phase 4: Extract Enterprise (Week 4-5)
|
||||
- Move enterprise_integration.py, compliance_engine.py, certification related -> enterprise
|
||||
- Create enterprise systemd service
|
||||
|
||||
### Phase 5: Extract Identity (Week 5-6)
|
||||
- Move auth/identity service files -> identity
|
||||
- Create identity systemd service
|
||||
|
||||
### Phase 6: Extract AI Models (Week 6-7)
|
||||
- Move advanced_*.py, multi_modal_fusion, ai verification -> ai-models
|
||||
- Create ai-models systemd service
|
||||
|
||||
### Phase 7: Extract Computing & Payment (Week 7-8)
|
||||
- Move gpu, resource, payment services to their own packages
|
||||
|
||||
### Phase 8: Final Integration (Week 8-9)
|
||||
- Update all clients to use new service endpoints
|
||||
- Test inter-service communication
|
||||
- Update documentation
|
||||
- Deprecate old monolith
|
||||
|
||||
## Files to Create/Modify
|
||||
|
||||
### New shared core (apps/coordinator-api/src/app/core/)
|
||||
- config.py (extracted from existing config.py)
|
||||
- logging.py (centralized logger setup)
|
||||
- database.py (SQLAlchemy session, Base)
|
||||
- exceptions.py (common exceptions)
|
||||
- security.py (auth dependencies)
|
||||
|
||||
### New service apps (47 directories total)
|
||||
Each: apps/<service>/src/app/{routers,services,models,main.py}
|
||||
|
||||
### Modified files
|
||||
- Root pyproject.toml: add service packages
|
||||
- Systemd: add 7 new .service files
|
||||
- Nginx config: new upstream blocks
|
||||
- Docker compose: add 7 new containers
|
||||
- Monitoring: new service endpoints for health
|
||||
|
||||
## Rollback Plan
|
||||
- Keep original monolith running alongside new services during transition
|
||||
- Use feature flags to route traffic
|
||||
- Comprehensive integration tests before cutover
|
||||
|
||||
## Success Criteria
|
||||
- Each service < 3,000 LOC (target 1,500)
|
||||
- Each service independently deployable
|
||||
- API contracts stable and documented
|
||||
- CI/CD per service
|
||||
239
.hermes/plans/2026-05-12_142930-agent-management-extraction.md
Normal file
239
.hermes/plans/2026-05-12_142930-agent-management-extraction.md
Normal file
@@ -0,0 +1,239 @@
|
||||
# Agent-Management Service Extraction Plan
|
||||
|
||||
## Overview
|
||||
|
||||
Extract the agent-related functionality from the coordinator-api monolith into a standalone microservice while maintaining operational continuity.
|
||||
|
||||
## Current State
|
||||
|
||||
**Monolith:** `apps/coordinator-api/src/app/`
|
||||
- Services: 46,594 LOC across 89 files
|
||||
- Domain layer: `domain/` contains all business entities (Agent, AgentExecution, AgentStatus, etc.)
|
||||
- Target agent files to extract: **18 files** (6 routers, 12 services)
|
||||
- Largest files: agent_service.py (1,159 LOC), agent_integration.py (1,117 LOC), agent_communication.py (988 LOC)
|
||||
|
||||
## Bounded Context: Agent-Management
|
||||
|
||||
**Responsibility:** AI agent lifecycle, orchestration, performance tracking, security, and marketplace registry.
|
||||
|
||||
**In-Scope Files:**
|
||||
|
||||
### Services (12)
|
||||
```
|
||||
services/agent_service.py (1,159 LOC)
|
||||
services/agent_integration.py (1,117 LOC)
|
||||
services/agent_communication.py (988 LOC)
|
||||
services/agent_orchestrator.py
|
||||
services/agent_performance_service.py
|
||||
services/agent_security.py
|
||||
services/agent_portfolio_manager.py
|
||||
services/agent_service_marketplace.py
|
||||
services/advanced_rl/agents.py (+ sub-agents: ppo_agent.py, rainbow_dqn_agent.py, sac_agent.py)
|
||||
```
|
||||
|
||||
### Routers (6)
|
||||
```
|
||||
routers/agent_router.py
|
||||
routers/agent_integration_router.py
|
||||
routers/agent_performance.py
|
||||
routers/agent_creativity.py
|
||||
routers/agent_security_router.py
|
||||
routers/services.py (agent services listing endpoint)
|
||||
```
|
||||
|
||||
## Critical Dependencies
|
||||
|
||||
1. **Domain Layer** (`app.domain`)
|
||||
- All agent services import from `..domain.agent` (AgentExecution, AgentStatus, AIAgentWorkflow, etc.)
|
||||
- Solution: Keep domain/ in monolith for now; new service imports via a **shared-domain package** to be created
|
||||
- Create `apps/shared-domain/src/app/domain/` as a symlink or copy that both services can import
|
||||
- Long-term: Extract entire domain layer to shared-domain package
|
||||
|
||||
2. **aitbc package**
|
||||
- Already available as root package. Use directly.
|
||||
|
||||
3. **SQLModel/SQLAlchemy**
|
||||
- Already in dependencies via root pyproject.toml
|
||||
|
||||
4. **Other monolith services**
|
||||
- Some routers may call agent endpoints. These will need to be updated to use HTTP client to new service (Phase 3 internal routing via nginx)
|
||||
|
||||
## Implementation Steps
|
||||
|
||||
### Step 0: Prepare Shared Domain Package (Prerequisite)
|
||||
- Create `apps/shared-domain/src/app/domain/`
|
||||
- Copy all files from coordinator-api's `domain/` EXCEPT non-agent ones if desired
|
||||
- Or simpler: symlink entire domain directory: `ln -s ../../coordinator-api/src/app/domain apps/shared-domain/src/app/`
|
||||
- Update imports in new service to use `from shared-domain.app.domain.agent import ...`
|
||||
- Add `shared-domain` to pyproject.toml dependencies in consuming services
|
||||
|
||||
**Recommendation:** Use symlink for rapid iteration, then formalize package later.
|
||||
|
||||
### Step 1: Create agent-management Service Skeleton
|
||||
```
|
||||
apps/agent-management/
|
||||
├── pyproject.toml
|
||||
├── README.md
|
||||
└── src/
|
||||
└── app/
|
||||
├── __init__.py
|
||||
├── main.py
|
||||
├── core/
|
||||
│ ├── __init__.py
|
||||
│ ├── config.py (import from shared-core)
|
||||
│ ├── logging.py (import from shared-core)
|
||||
│ └── database.py (import from shared-core)
|
||||
├── domain/ → symlink to ../../shared-domain/src/app/domain
|
||||
├── routers/
|
||||
│ ├── __init__.py
|
||||
│ ├── agent_router.py (copied & adapted)
|
||||
│ ├── agent_integration_router.py
|
||||
│ ├── agent_performance.py
|
||||
│ ├── agent_creativity.py
|
||||
│ ├── agent_security_router.py
|
||||
│ └── services.py
|
||||
└── services/
|
||||
├── __init__.py
|
||||
├── agent_service.py
|
||||
├── agent_orchestrator.py
|
||||
├── agent_communication.py
|
||||
├── agent_performance_service.py
|
||||
├── agent_security.py
|
||||
├── agent_integration.py
|
||||
├── agent_portfolio_manager.py
|
||||
├── agent_service_marketplace.py
|
||||
└── advanced_rl/
|
||||
├── __init__.py
|
||||
├── agents.py
|
||||
└── ppo_agent.py, rainbow_dqn_agent.py, sac_agent.py
|
||||
```
|
||||
|
||||
### Step 2: Adapt Code for Service Boundaries
|
||||
|
||||
**Changes needed per file:**
|
||||
|
||||
- Update all `from ..domain.agent import X` to `from shared-domain.app.domain.agent import X`
|
||||
- Remove any imports from other monolith services (e.g., `from ..services.other_service import X`)
|
||||
- Replace internal service calls with HTTP client calls or event bus (defer to later phase)
|
||||
- Update `ServiceSettings` to use agent-management specific defaults (port 8012)
|
||||
- Add health check endpoint (already in template)
|
||||
- Verify database setup: AgentExecution etc use shared Base. Need to call `Base.metadata.create_all(bind=engine)` on startup
|
||||
|
||||
**Special Case: advanced_rl/**
|
||||
- These are AI model inference services. Consider moving to `ai-models` service instead.
|
||||
- For now, keep in agent-management to maintain functionality.
|
||||
|
||||
### Step 3: Update Monolith to Proxy Requests (During Transition)
|
||||
|
||||
**Option A: Nginx Routing**
|
||||
- Add nginx upstream for agent-management on port 8012
|
||||
- Change coordinator-api routes for `/api/v1/agent/*` to proxy to agent-management
|
||||
- Monolith no longer handles agent endpoints
|
||||
|
||||
**Option B: In-app Redirection**
|
||||
- Keep routers in monolith but replace handlers with `HTTPClient` calls to new service
|
||||
- More gradual migration but adds latency
|
||||
|
||||
**Recommendation:** Option A - cleaner separation, easier to rollback.
|
||||
|
||||
### Step 4: Create Systemd Service
|
||||
|
||||
```
|
||||
/etc/systemd/system/aitbc-agent-management.service
|
||||
[Unit]
|
||||
Description=AITBC Agent Management Service
|
||||
After=network.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
User=aitbc
|
||||
WorkingDirectory=/opt/aitbc/apps/agent-management
|
||||
Environment=PATH=/opt/aitbc/venv/bin
|
||||
Environment=PYTHONPATH=/opt/aitbc
|
||||
ExecStart=/opt/aitbc/venv/bin/uvicorn app.main:app --host 127.0.0.1 --port 8012
|
||||
Restart=on-failure
|
||||
RestartSec=10
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
```
|
||||
|
||||
### Step 5: Database Migration
|
||||
|
||||
- Agent domain models likely already have tables defined via SQLModel
|
||||
- In `main.py` startup event, call `Base.metadata.create_all(bind=engine)` to ensure tables exist
|
||||
- Ensure the new service uses same database as monolith (coordinator.db) initially
|
||||
- Later: separate database (Phase 8)
|
||||
|
||||
### Step 6: Integration Testing
|
||||
|
||||
1. Start agent-management service
|
||||
2. Verify health endpoint: `curl http://localhost:8012/health`
|
||||
3. Test agent creation via API
|
||||
4. Verify coordinator-api can still access agent data (through new service or direct DB if keeping shared DB)
|
||||
5. Run existing integration tests against new service
|
||||
|
||||
### Step 7: Update Coordinator-API
|
||||
|
||||
- Remove the 18 extracted files from monolith
|
||||
- Remove domain/agent related imports from remaining monolith services if they now use agent-management API
|
||||
- Update any remaining references to agent endpoints to use HTTP client or nginx proxy
|
||||
|
||||
### Step 8: Documentation & Monitoring
|
||||
|
||||
- Update README with agent-management API docs
|
||||
- Add metrics endpoint if enabled
|
||||
- Update deployment scripts
|
||||
|
||||
## Rollback Plan
|
||||
|
||||
1. Keep monolith files in git history (do not delete, just move)
|
||||
2. Keep nginx config either/or - can revert upstream routing
|
||||
3. Database shared initially, so data is accessible to both
|
||||
4. Systemd service can be disabled; monolith still runs
|
||||
|
||||
## Success Criteria
|
||||
|
||||
- [ ] Agent-management service starts and health check passes on port 8012
|
||||
- [ ] Can create/query agents via API
|
||||
- [ ] Existing coordinator-api functionality that depends on agents still works
|
||||
- [ ] No errors in logs during integration test
|
||||
- [ ] Systemd service auto-restarts on failure
|
||||
|
||||
## Open Questions
|
||||
|
||||
1. **RL Agents**: Should advanced_rl be part of agent-management or ai-models?
|
||||
- Recommendation: Keep in agent-management for now (AI agent inference is part of agent runtime). Can split later if ai-models becomes a separate inference service.
|
||||
|
||||
2. **Database**: Separate or shared?
|
||||
- Phase 1: Shared (same coordinator.db) for simplicity
|
||||
- Phase 8: Split to dedicated agent-management database
|
||||
|
||||
3. **Cross-service calls**: Currently agent integration uses other services directly (imports). Need to replace with HTTP or event bus.
|
||||
- Defer until Phase 8 (Final Integration) to avoid breaking existing flow
|
||||
|
||||
4. **Domain extraction**: The domain models are currently in monolith. Should we extract entire domain to a package?
|
||||
- Immediate need: Create shared-domain package (symlink) to break import cycle
|
||||
- Future: Extract domain to true package with independent version
|
||||
|
||||
## Timeline Estimate
|
||||
|
||||
- Step 0 (shared-domain): 2h
|
||||
- Step 1 (skeleton): 4h
|
||||
- Step 2 (adaptation): 8h (bulk of work - fixing imports, resolving dependencies)
|
||||
- Step 3 (nginx routing): 2h
|
||||
- Step 4 (systemd): 1h
|
||||
- Step 5 (DB): 1h
|
||||
- Step 6 (testing): 4h
|
||||
- Step 7 (monolith cleanup): 4h
|
||||
- Step 8 (docs): 2h
|
||||
|
||||
**Total: ~28 hours (3-4 days)**
|
||||
|
||||
## Risks
|
||||
|
||||
- Hidden dependencies on other monolith services may cause runtime import errors
|
||||
- Domain models may have cross-references that require co-migration
|
||||
- Database migrations may be needed if agent tables don't exist yet
|
||||
- Existing integration tests may fail and need updating
|
||||
- Breaking changes if API contracts differ from original
|
||||
218
.hermes/plans/2026-05-12_150000-tighten-mypy-config.md
Normal file
218
.hermes/plans/2026-05-12_150000-tighten-mypy-config.md
Normal file
@@ -0,0 +1,218 @@
|
||||
|
||||
# Tighten Mypy Configuration Plan
|
||||
|
||||
## Current State
|
||||
|
||||
**Root pyproject.toml [tool.mypy] settings:**
|
||||
```toml
|
||||
warn_return_any = true
|
||||
warn_unused_configs = true
|
||||
check_untyped_defs = false
|
||||
disallow_incomplete_defs = false
|
||||
disallow_untyped_defs = false
|
||||
disallow_untyped_decorators = false
|
||||
no_implicit_optional = false
|
||||
warn_redundant_casts = false
|
||||
warn_unused_ignores = false
|
||||
warn_no_return = true
|
||||
warn_unreachable = false
|
||||
strict_equality = false
|
||||
```
|
||||
|
||||
**Overrides:**
|
||||
- Heavy libraries (torch, cv2, pandas, numpy, web3, etc.) are `ignore_missing_imports = true`
|
||||
- Coordiator-api modules are `ignore_errors = true` (catch-all)
|
||||
|
||||
This is **extremely permissive** - essentially just warns on return_any and missing configs. It does not enforce:
|
||||
- Function argument/return type completeness
|
||||
- Avoiding implicit `Any`
|
||||
- Avoiding unnecessary type: ignore comments
|
||||
- Detecting unreachable code
|
||||
- Strict equality checks (None vs False)
|
||||
|
||||
## Proposed Tightening Phases
|
||||
|
||||
### Phase 1: Enable Foundational Checks (Low Effort, High Value)
|
||||
Target: enable 4 key options that catch real bugs with minimal friction
|
||||
|
||||
```toml
|
||||
disallow_untyped_defs = true
|
||||
disallow_incomplete_defs = true
|
||||
warn_redundant_casts = true
|
||||
warn_unused_ignores = true
|
||||
```
|
||||
|
||||
**Impact:**
|
||||
- Functions must have complete type signatures (all args+returns typed)
|
||||
- Redundant cast() calls will be flagged
|
||||
- Unused `# type: ignore` comments will be flagged
|
||||
- Minimal code changes required (most functions already typed)
|
||||
|
||||
**Estimated effort:**
|
||||
- 1 hour to update config
|
||||
- 2-4 hours to fix violations in production code
|
||||
- Total: ~1 day
|
||||
|
||||
**Validation:**
|
||||
- Run `mypy apps` and ensure 0 errors
|
||||
- Keep existing overrides for external libraries and coordinator-api
|
||||
|
||||
### Phase 2: Stricter Optional Handling (Medium Effort)
|
||||
Enable:
|
||||
```toml
|
||||
no_implicit_optional = true
|
||||
warn_unreachable = true
|
||||
strict_equality = true
|
||||
```
|
||||
|
||||
**Impact:**
|
||||
- Variables defaulting to `None` must be explicitly `Optional[...]`
|
||||
- Unreachable code will be flagged (dead code detection)
|
||||
- Equality comparisons with None must use `is` not `==`
|
||||
|
||||
**Estimated effort:** 2-3 days to fix violations across codebase
|
||||
|
||||
### Phase 3: Gradual Per-Module Strictness (Long-term)
|
||||
- Move coordinator-api out of catch-all `ignore_errors`
|
||||
- Add per-module overrides as we achieve correctness
|
||||
- Eventually remove `ignore_errors` blanket
|
||||
|
||||
**Estimated effort:** Ongoing as part of decomposition
|
||||
|
||||
## Implementation Steps
|
||||
|
||||
### Step 1: Backup Current Config
|
||||
```bash
|
||||
cp pyproject.toml pyproject.toml.backup
|
||||
```
|
||||
|
||||
### Step 2: Update Root Configuration
|
||||
|
||||
Modify `/opt/aitbc/pyproject.toml` [tool.mypy] section:
|
||||
|
||||
```diff
|
||||
[tool.mypy]
|
||||
python_version = "3.13"
|
||||
warn_return_any = true
|
||||
warn_unused_configs = true
|
||||
check_untyped_defs = false
|
||||
-disallow_incomplete_defs = false
|
||||
-disallow_untyped_defs = false
|
||||
+disallow_incomplete_defs = true
|
||||
+disallow_untyped_defs = true
|
||||
disallow_untyped_decorators = false
|
||||
no_implicit_optional = false
|
||||
warn_redundant_casts = false
|
||||
warn_unused_ignores = false
|
||||
warn_no_return = true
|
||||
warn_unreachable = false
|
||||
strict_equality = false
|
||||
```
|
||||
|
||||
### Step 3: Run Mypy and Collect Errors
|
||||
|
||||
```bash
|
||||
cd /opt/aitbc
|
||||
venv/bin/mypy apps --show-error-codes --no-color-output > mypy_errors.txt 2>&1
|
||||
```
|
||||
|
||||
### Step 4: Categorize Errors
|
||||
|
||||
Typical violations we'll see:
|
||||
- `Function is missing a return type annotation` (from disallow_untyped_defs)
|
||||
- `Function is missing a type annotation for one or more arguments` (from disallow_untyped_defs)
|
||||
- `Class is missing type parameters for generic type` (rare)
|
||||
- `dict, list, etc. used without type parameters` (from disallow_incomplete_defs)
|
||||
- `Redundant cast to X` (from warn_redundant_casts)
|
||||
- `Unused "type: ignore" comment` (from warn_unused_ignores)
|
||||
|
||||
### Step 5: Fix in Order of Impact
|
||||
|
||||
**A. Add missing type annotations to functions**
|
||||
- Priority: functions in shared-core, services, routers
|
||||
- Use explicit return types; if truly dynamic, use `-> Any` (but rarely needed)
|
||||
- Example:
|
||||
```python
|
||||
def get_engine(settings): # BEFORE
|
||||
def get_engine(settings: ServiceSettings) -> Engine: # AFTER
|
||||
```
|
||||
|
||||
**B. Add generic type parameters**
|
||||
- `list` -> `List[str]` or `list[int]`
|
||||
- `dict` -> `Dict[str, Any]`
|
||||
- Use `from typing import List, Dict`
|
||||
|
||||
**C. Remove redundant casts**
|
||||
- Delete `cast(Type, value)` if type is already clear to mypy
|
||||
- Use `reveal_type(value)` to check actual inferred type before removing
|
||||
|
||||
**D. Remove unused type: ignore**
|
||||
- Some `# type: ignore` comments are legacy and no longer needed
|
||||
- Delete them; if mypy still fails, leave or fix underlying issue
|
||||
|
||||
### Step 6: Iterate and Validate
|
||||
|
||||
After fixing categories, re-run mypy. Continue until `mypy apps` exits with code 0.
|
||||
|
||||
**Note:** We preserve `ignore_missing_imports` for heavy libraries, and `ignore_errors` for coordinator-api (since we're deferring decomposition).
|
||||
|
||||
### Step 7: Add CI Enforcement
|
||||
|
||||
Update pre-commit hooks or CI to run mypy on PRs:
|
||||
```yaml
|
||||
# .pre-commit-config.yaml or GitHub Actions
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: mypy
|
||||
name: mypy
|
||||
entry: mypy apps
|
||||
language: system
|
||||
pass_filenames: false
|
||||
```
|
||||
|
||||
## Rollback Plan
|
||||
|
||||
If the effort becomes too large:
|
||||
1. Revert pyproject.toml from backup
|
||||
2. Keep per-module `# mypy: ignore-errors` as needed
|
||||
3. Approach incrementally: enable one flag at a time
|
||||
|
||||
## Success Criteria
|
||||
|
||||
- `mypy apps` completes with 0 errors
|
||||
- No new type: ignore comments added without explanation
|
||||
- Production code has complete type signatures
|
||||
- CI pipeline includes mypy check
|
||||
|
||||
## Risks & Mitigations
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|------------|
|
||||
| Overwhelming number of errors | Enable flags incrementally (2 at a time), fix in batches by module |
|
||||
| Breaking existing functionality by incorrect type fixes | Run test suite after each batch; use `reveal_type` to debug |
|
||||
| Third-party library types incompatible | Keep `ignore_missing_imports` for those packages |
|
||||
| Coordinator-api too messy to fix now | Keep `ignore_errors` override; revisit after decomposition |
|
||||
|
||||
## Related Tasks
|
||||
|
||||
- **Decompose coordinator-api** - Once strict mypy is in place, easier to validate new services
|
||||
- **Shared-core library** - Strict typing ensures compatibility across services
|
||||
- **Connection pooling** - Use proper typed database sessions
|
||||
|
||||
## Open Questions
|
||||
|
||||
1. Should we also enable `strict` mode for new services? (Probably yes)
|
||||
2. Should we add type-checking to pre-commit hook for changed files only? (Yes, use `mypy --files <changed>`)
|
||||
3. How to handle legacy coordinator-api code? (Keep ignore_errors for now)
|
||||
|
||||
## Estimated Timeline
|
||||
|
||||
- **0-2 days:** Implement Phase 1, fix immediate violations
|
||||
- **3-7 days:** Address accumulated type errors, reach clean mypy
|
||||
- **Week 2:** Add CI enforcement, document guidelines
|
||||
- **Ongoing:** Maintain strict typing in new code
|
||||
|
||||
## References
|
||||
|
||||
- Mypy configuration: https://mypy.readthedocs.io/en/stable/config_file.html
|
||||
- Strict mode: https://mypy.readthedocs.io/en/stable/command_line.html#cmdoption-mypy-strict
|
||||
@@ -193,7 +193,7 @@ class Web3Client:
|
||||
})
|
||||
if len(transactions) >= limit:
|
||||
break
|
||||
except:
|
||||
except (KeyError, ValueError, AttributeError):
|
||||
continue
|
||||
|
||||
return transactions
|
||||
|
||||
@@ -206,7 +206,7 @@ class TestHelpers:
|
||||
try:
|
||||
os.remove(file_path)
|
||||
count += 1
|
||||
except:
|
||||
except (OSError, IOError):
|
||||
pass
|
||||
return count
|
||||
|
||||
@@ -389,7 +389,7 @@ import time
|
||||
def create_test_scenario(name: str, steps: List[Callable]) -> Callable:
|
||||
"""Create a test scenario with multiple steps"""
|
||||
def scenario():
|
||||
print(f"Running test scenario: {name}")
|
||||
logger.info("Running test scenario", name=name)
|
||||
results = []
|
||||
for i, step in enumerate(steps):
|
||||
try:
|
||||
|
||||
@@ -324,7 +324,7 @@ jwt_secret = os.getenv("JWT_SECRET")
|
||||
if not jwt_secret:
|
||||
raise ValueError(
|
||||
"JWT_SECRET environment variable must be set. "
|
||||
"Generate a secure secret using: python -c 'import secrets; print(secrets.token_urlsafe(32))'"
|
||||
"Generate a secure secret using: python -c 'import secrets; logger.info(secrets.token_urlsafe(32))'"
|
||||
)
|
||||
jwt_handler = JWTHandler(jwt_secret)
|
||||
password_manager = PasswordManager()
|
||||
|
||||
@@ -74,7 +74,7 @@ class Settings(BaseSettings):
|
||||
connection_timeout: int = 30
|
||||
|
||||
# Security settings
|
||||
secret_key: str = "your-secret-key-change-in-production"
|
||||
secret_key: str
|
||||
allowed_hosts: list = ["*"]
|
||||
cors_origins: list = ["*"]
|
||||
|
||||
@@ -237,7 +237,7 @@ class EnvironmentConfig:
|
||||
"enable_metrics": True,
|
||||
"workers": 4,
|
||||
"cors_origins": ["https://aitbc.com"],
|
||||
"secret_key": os.getenv("SECRET_KEY", "change-this-in-production"),
|
||||
"secret_key": os.getenv("SECRET_KEY"),
|
||||
"allowed_hosts": ["aitbc.com", "www.aitbc.com"]
|
||||
}
|
||||
|
||||
@@ -275,7 +275,7 @@ class ConfigLoader:
|
||||
errors = []
|
||||
|
||||
# Validate required settings
|
||||
if not settings.secret_key or settings.secret_key == "your-secret-key-change-in-production":
|
||||
if not settings.secret_key:
|
||||
if settings.environment == Environment.PRODUCTION:
|
||||
errors.append("SECRET_KEY must be set in production")
|
||||
|
||||
|
||||
@@ -39,11 +39,15 @@ async def login(login_data: Dict[str, str]):
|
||||
import os
|
||||
|
||||
demo_users = {
|
||||
"admin": os.getenv("DEMO_ADMIN_PASSWORD", "admin123"),
|
||||
"operator": os.getenv("DEMO_OPERATOR_PASSWORD", "operator123"),
|
||||
"user": os.getenv("DEMO_USER_PASSWORD", "user123")
|
||||
"admin": os.getenv("DEMO_ADMIN_PASSWORD"),
|
||||
"operator": os.getenv("DEMO_OPERATOR_PASSWORD"),
|
||||
"user": os.getenv("DEMO_USER_PASSWORD")
|
||||
}
|
||||
|
||||
# Require environment variables for demo credentials - no hardcoded fallbacks
|
||||
if username in demo_users and demo_users[username] is None:
|
||||
raise HTTPException(status_code=500, detail=f"{username.capitalize()} password not configured in environment")
|
||||
|
||||
if username == "admin" and password == demo_users["admin"]:
|
||||
user_id = "admin_001"
|
||||
role = Role.ADMIN
|
||||
|
||||
@@ -80,7 +80,7 @@ async def send_message(request: MessageRequest):
|
||||
if state.communication_manager:
|
||||
try:
|
||||
await state.communication_manager.send_message(protocol, message)
|
||||
except:
|
||||
except Exception:
|
||||
pass # Protocol send is optional
|
||||
|
||||
return {
|
||||
@@ -172,7 +172,7 @@ async def broadcast_message(request: BroadcastRequest):
|
||||
if state.communication_manager:
|
||||
try:
|
||||
await state.communication_manager.send_message("broadcast", message)
|
||||
except:
|
||||
except Exception:
|
||||
pass # Protocol send is optional
|
||||
|
||||
return {
|
||||
|
||||
@@ -629,7 +629,7 @@ async def example_usage():
|
||||
"status": "active"
|
||||
})
|
||||
|
||||
print(f"Found {len(agents)} agents")
|
||||
logger.info(f"Found {len(agents)} agents")
|
||||
|
||||
# Find best agent
|
||||
best_agent = await discovery_service.find_best_agent({
|
||||
@@ -638,7 +638,7 @@ async def example_usage():
|
||||
})
|
||||
|
||||
if best_agent:
|
||||
print(f"Best agent: {best_agent.agent_id}")
|
||||
logger.info(f"Best agent: {best_agent.agent_id}")
|
||||
|
||||
await registry.stop()
|
||||
|
||||
|
||||
@@ -55,7 +55,7 @@ class MessageStorage:
|
||||
# Try to parse ISO format
|
||||
dt = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
|
||||
timestamp_float = dt.timestamp()
|
||||
except:
|
||||
except Exception:
|
||||
# Already a float or int
|
||||
timestamp_float = float(timestamp_str)
|
||||
await self.redis.zadd(f"messages:timestamp", {message_id: timestamp_float})
|
||||
|
||||
26
apps/agent-management/pyproject.toml
Normal file
26
apps/agent-management/pyproject.toml
Normal file
@@ -0,0 +1,26 @@
|
||||
[tool.poetry]
|
||||
name = "aitbc-agent-management"
|
||||
version = "0.1.0"
|
||||
description = "AITBC Agent Management Service - AI agent lifecycle, orchestration, and performance tracking"
|
||||
authors = ["AITBC Team <team@aitbc.dev>"]
|
||||
readme = "README.md"
|
||||
packages = [{include = "app", from = "src"}]
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.13"
|
||||
aitbc = {path = "../../../"}
|
||||
aitbc-shared-domain = {path = "../../shared-domain"}
|
||||
aitbc-shared-core = {path = "../../shared-core"}
|
||||
fastapi = ">=0.104.0"
|
||||
uvicorn = ">=0.24.0"
|
||||
sqlmodel = ">=0.0.14"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
pytest = ">=9.0.3"
|
||||
pytest-asyncio = ">=1.3.0"
|
||||
pytest-cov = ">=6.0.0"
|
||||
httpx = ">=0.28.1"
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
0
apps/agent-management/src/app/__init__.py
Normal file
0
apps/agent-management/src/app/__init__.py
Normal file
0
apps/agent-management/src/app/core/__init__.py
Normal file
0
apps/agent-management/src/app/core/__init__.py
Normal file
70
apps/agent-management/src/app/core/config.py
Normal file
70
apps/agent-management/src/app/core/config.py
Normal file
@@ -0,0 +1,70 @@
|
||||
"""Configuration for Agent Management Service"""
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class DatabaseConfig(BaseSettings):
|
||||
"""Database configuration with adapter selection."""
|
||||
|
||||
adapter: str = "sqlite" # sqlite, postgresql
|
||||
url: Optional[str] = None
|
||||
pool_size: int = 10
|
||||
max_overflow: int = 20
|
||||
pool_pre_ping: bool = True
|
||||
|
||||
@property
|
||||
def effective_url(self) -> str:
|
||||
"""Get the effective database URL."""
|
||||
if self.url:
|
||||
return self.url
|
||||
if self.adapter == "sqlite":
|
||||
# Use absolute path from DATA_DIR if available
|
||||
import os
|
||||
data_dir = os.getenv("DATA_DIR", "/opt/aitbc/data")
|
||||
return f"sqlite:///{data_dir}/coordinator.db"
|
||||
return f"{self.adapter}://localhost:5432/agent_management"
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env", env_file_encoding="utf-8", case_sensitive=False, extra="allow"
|
||||
)
|
||||
|
||||
|
||||
class ServiceSettings(BaseSettings):
|
||||
"""Base settings for AITBC microservices."""
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env", env_file_encoding="utf-8", case_sensitive=False, extra="allow"
|
||||
)
|
||||
|
||||
# Environment
|
||||
service_name: str = "aitbc-service"
|
||||
app_env: str = "dev"
|
||||
app_host: str = "127.0.0.1"
|
||||
app_port: int = 8000
|
||||
debug: bool = False
|
||||
|
||||
# Logging
|
||||
log_level: str = "INFO"
|
||||
log_dir: str = "/var/log/aitbc/services"
|
||||
|
||||
# Database
|
||||
database: DatabaseConfig = DatabaseConfig()
|
||||
|
||||
# API
|
||||
api_prefix: str = "/api/v1"
|
||||
|
||||
# Feature flags
|
||||
enable_metrics: bool = True
|
||||
enable_health_check: bool = True
|
||||
|
||||
# API Keys (comma-separated in env)
|
||||
admin_api_keys: List[str] = Field(default_factory=list)
|
||||
client_api_keys: List[str] = Field(default_factory=list)
|
||||
miner_api_keys: List[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
# Global settings instance
|
||||
settings = ServiceSettings()
|
||||
36
apps/agent-management/src/app/core/database.py
Normal file
36
apps/agent-management/src/app/core/database.py
Normal file
@@ -0,0 +1,36 @@
|
||||
"""Shared database utilities for AITBC services."""
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker, declarative_base
|
||||
from typing import Generator
|
||||
|
||||
from .config import ServiceSettings
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
def get_engine(settings: ServiceSettings):
|
||||
"""Create SQLAlchemy engine based on configuration."""
|
||||
db_config = settings.database
|
||||
return create_engine(
|
||||
db_config.effective_url,
|
||||
pool_size=db_config.pool_size,
|
||||
max_overflow=db_config.max_overflow,
|
||||
pool_pre_ping=db_config.pool_pre_ping,
|
||||
echo=settings.debug
|
||||
)
|
||||
|
||||
|
||||
def get_sessionmaker(engine):
|
||||
"""Create session factory."""
|
||||
return sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
||||
|
||||
|
||||
def get_db(engine) -> Generator:
|
||||
"""Dependency for FastAPI endpoints."""
|
||||
Session = get_sessionmaker(engine)
|
||||
db = Session()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
66
apps/agent-management/src/app/core/logging.py
Normal file
66
apps/agent-management/src/app/core/logging.py
Normal file
@@ -0,0 +1,66 @@
|
||||
"""Shared logging configuration for AITBC services."""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from ..core.config import ServiceSettings
|
||||
|
||||
|
||||
def setup_logging(settings: Optional[ServiceSettings] = None, level: str = None) -> logging.Logger:
|
||||
"""Configure structured logging for the service.
|
||||
|
||||
Args:
|
||||
settings: Service settings containing log configuration
|
||||
level: Override log level
|
||||
|
||||
Returns:
|
||||
Configured root logger
|
||||
"""
|
||||
if settings:
|
||||
log_level = level or settings.log_level
|
||||
log_dir = Path(settings.log_dir)
|
||||
else:
|
||||
log_level = level or "INFO"
|
||||
log_dir = Path("/var/log/aitbc/services")
|
||||
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Create formatter
|
||||
formatter = logging.Formatter(
|
||||
fmt="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S"
|
||||
)
|
||||
|
||||
# Configure root logger
|
||||
root_logger = logging.getLogger()
|
||||
root_logger.setLevel(getattr(logging, log_level.upper()))
|
||||
|
||||
# Clear existing handlers
|
||||
root_logger.handlers.clear()
|
||||
|
||||
# Console handler
|
||||
console_handler = logging.StreamHandler(sys.stdout)
|
||||
console_handler.setFormatter(formatter)
|
||||
root_logger.addHandler(console_handler)
|
||||
|
||||
# File handler
|
||||
if settings and settings.service_name:
|
||||
file_handler = logging.FileHandler(
|
||||
log_dir / f"{settings.service_name}.log"
|
||||
)
|
||||
file_handler.setFormatter(formatter)
|
||||
root_logger.addHandler(file_handler)
|
||||
|
||||
return root_logger
|
||||
|
||||
|
||||
def get_logger(name: str) -> logging.Logger:
|
||||
"""Get a logger with the given name.
|
||||
|
||||
Usage:
|
||||
from app.core.logging import get_logger
|
||||
logger = get_logger(__name__)
|
||||
"""
|
||||
return logging.getLogger(name)
|
||||
73
apps/agent-management/src/app/deps.py
Executable file
73
apps/agent-management/src/app/deps.py
Executable file
@@ -0,0 +1,73 @@
|
||||
"""Dependency injection module for AITBC Agent Management Service
|
||||
|
||||
Provides unified dependency injection using ServiceSettings.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from fastapi import Header, HTTPException
|
||||
|
||||
from .core.config import settings # We'll create this file
|
||||
|
||||
|
||||
def _validate_api_key(allowed_keys: list[str], api_key: str | None) -> str:
|
||||
# In development mode, allow any API key for testing
|
||||
import os
|
||||
|
||||
if os.getenv("APP_ENV", "dev") == "dev":
|
||||
return api_key or "dev_key"
|
||||
|
||||
allowed = {key.strip() for key in allowed_keys if key}
|
||||
if not api_key or api_key not in allowed:
|
||||
raise HTTPException(status_code=401, detail="invalid api key")
|
||||
return api_key
|
||||
|
||||
|
||||
def require_client_key() -> Callable[[str | None], str]:
|
||||
"""Dependency for client API key authentication (reads live settings)."""
|
||||
|
||||
def validator(api_key: str | None = Header(default=None, alias="X-Api-Key")) -> str:
|
||||
return _validate_api_key(settings.client_api_keys, api_key)
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
def require_miner_key() -> Callable[[str | None], str]:
|
||||
"""Dependency for miner API key authentication (reads live settings)."""
|
||||
|
||||
def validator(api_key: str | None = Header(default=None, alias="X-Api-Key")) -> str:
|
||||
return _validate_api_key(settings.miner_api_keys, api_key)
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
def get_miner_id() -> Callable[[str | None], str]:
|
||||
"""Dependency to get miner ID from X-Miner-ID header."""
|
||||
|
||||
def validator(miner_id: str | None = Header(default=None, alias="X-Miner-ID")) -> str:
|
||||
if not miner_id:
|
||||
raise HTTPException(status_code=400, detail="X-Miner-ID header required")
|
||||
return miner_id
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
def require_admin_key() -> Callable[[str | None], str]:
|
||||
"""Dependency for admin API key authentication (reads live settings)."""
|
||||
|
||||
def validator(api_key: str | None = Header(default=None, alias="X-Api-Key")) -> str:
|
||||
return _validate_api_key(settings.admin_api_keys, api_key)
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
# Legacy APIKeyValidator class for backward compatibility with tests
|
||||
class APIKeyValidator:
|
||||
"""Legacy API key validator class for backward compatibility."""
|
||||
|
||||
def __init__(self, allowed_keys: list[str]):
|
||||
self.allowed_keys = allowed_keys
|
||||
|
||||
def __call__(self, api_key: str | None = None) -> str:
|
||||
"""Validate API key."""
|
||||
return _validate_api_key(self.allowed_keys, api_key)
|
||||
1
apps/agent-management/src/app/domain
Symbolic link
1
apps/agent-management/src/app/domain
Symbolic link
@@ -0,0 +1 @@
|
||||
../../coordinator-api/src/app/domain
|
||||
83
apps/agent-management/src/app/main.py
Normal file
83
apps/agent-management/src/app/main.py
Normal file
@@ -0,0 +1,83 @@
|
||||
#!/usr/bin/env python3
|
||||
"""AITBC Agent Management Service"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add project root to path
|
||||
project_root = Path(__file__).parent.parent.parent.parent.parent
|
||||
if str(project_root) not in sys.path:
|
||||
sys.path.insert(0, str(project_root))
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from aitbc import get_logger
|
||||
|
||||
# Local imports
|
||||
from .core.config import settings
|
||||
from .core.logging import setup_logging, get_logger
|
||||
from .core.database import Base, get_engine, get_sessionmaker
|
||||
|
||||
# Setup logging
|
||||
setup_logging(settings)
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI(
|
||||
title="AITBC Agent Management API",
|
||||
description="AI agent lifecycle, orchestration, performance tracking, and security",
|
||||
version="0.1.0",
|
||||
debug=settings.debug
|
||||
)
|
||||
|
||||
# Database setup
|
||||
engine = get_engine(settings)
|
||||
SessionLocal = get_sessionmaker(engine)
|
||||
|
||||
# Create tables on startup
|
||||
@app.on_event("startup")
|
||||
def on_startup():
|
||||
Base.metadata.create_all(bind=engine)
|
||||
logger.info("Agent Management service started")
|
||||
|
||||
# Dependency
|
||||
def get_db():
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
# Include routers
|
||||
from .routers import (
|
||||
agent_router,
|
||||
agent_integration_router,
|
||||
agent_performance,
|
||||
agent_creativity,
|
||||
agent_security_router,
|
||||
services as agent_services_router
|
||||
)
|
||||
|
||||
# Mount routers with prefix
|
||||
app.include_router(agent_router.router, prefix=f"{settings.api_prefix}/agents")
|
||||
app.include_router(agent_integration_router.router, prefix=f"{settings.api_prefix}/agents/integration")
|
||||
app.include_router(agent_performance.router, prefix=f"{settings.api_prefix}/agents/performance")
|
||||
app.include_router(agent_creativity.router, prefix=f"{settings.api_prefix}/agents/creativity")
|
||||
app.include_router(agent_security_router.router, prefix=f"{settings.api_prefix}/agents/security")
|
||||
app.include_router(agent_services_router.router, prefix=f"{settings.api_prefix}/services")
|
||||
|
||||
@app.get("/health")
|
||||
def health_check():
|
||||
return {"status": "healthy", "service": settings.service_name}
|
||||
|
||||
@app.get("/")
|
||||
def root():
|
||||
return {"message": "Welcome to AITBC Agent Management Service"}
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(
|
||||
"app.main:app",
|
||||
host=settings.app_host,
|
||||
port=settings.app_port,
|
||||
reload=settings.debug
|
||||
)
|
||||
0
apps/agent-management/src/app/models/__init__.py
Normal file
0
apps/agent-management/src/app/models/__init__.py
Normal file
17
apps/agent-management/src/app/routers/__init__.py
Normal file
17
apps/agent-management/src/app/routers/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""Agent Management Routers"""
|
||||
|
||||
from .agent_router import router as agent_router
|
||||
from .agent_integration_router import router as agent_integration_router
|
||||
from .agent_performance import router as agent_performance_router
|
||||
from .agent_creativity import router as agent_creativity_router
|
||||
from .agent_security_router import router as agent_security_router
|
||||
from .services import router as services_router
|
||||
|
||||
__all__ = [
|
||||
"agent_router",
|
||||
"agent_integration_router",
|
||||
"agent_performance_router",
|
||||
"agent_creativity_router",
|
||||
"agent_security_router",
|
||||
"services_router",
|
||||
]
|
||||
196
apps/agent-management/src/app/routers/agent_creativity.py
Executable file
196
apps/agent-management/src/app/routers/agent_creativity.py
Executable file
@@ -0,0 +1,196 @@
|
||||
from typing import Annotated
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
"""
|
||||
Agent Creativity API Endpoints
|
||||
REST API for agent creativity enhancement, ideation, and cross-domain synthesis
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from app.domain.agent_performance import CreativeCapability
|
||||
from sqlmodel import select
|
||||
from ..services.creative_capabilities_service import (
|
||||
CreativityEnhancementEngine,
|
||||
CrossDomainCreativeIntegrator,
|
||||
IdeationAlgorithm,
|
||||
)
|
||||
from ..storage import get_session
|
||||
|
||||
router = APIRouter(prefix="/v1/agent-creativity", tags=["agent-creativity"])
|
||||
|
||||
|
||||
# Models
|
||||
class CreativeCapabilityCreate(BaseModel):
|
||||
agent_id: str
|
||||
creative_domain: str = Field(..., description="e.g., artistic, design, innovation, scientific, narrative")
|
||||
capability_type: str = Field(..., description="e.g., generative, compositional, analytical, innovative")
|
||||
generation_models: list[str]
|
||||
initial_score: float = Field(0.5, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class CreativeCapabilityResponse(BaseModel):
|
||||
capability_id: str
|
||||
agent_id: str
|
||||
creative_domain: str
|
||||
capability_type: str
|
||||
originality_score: float
|
||||
novelty_score: float
|
||||
aesthetic_quality: float
|
||||
coherence_score: float
|
||||
style_variety: int
|
||||
creative_specializations: list[str]
|
||||
status: str
|
||||
|
||||
|
||||
class EnhanceCreativityRequest(BaseModel):
|
||||
algorithm: str = Field(
|
||||
"divergent_thinking",
|
||||
description="divergent_thinking, conceptual_blending, morphological_analysis, lateral_thinking, bisociation",
|
||||
)
|
||||
training_cycles: int = Field(100, ge=1, le=1000)
|
||||
|
||||
|
||||
class EvaluateCreationRequest(BaseModel):
|
||||
creation_data: dict[str, Any]
|
||||
expert_feedback: dict[str, float] | None = None
|
||||
|
||||
|
||||
class IdeationRequest(BaseModel):
|
||||
problem_statement: str
|
||||
domain: str
|
||||
technique: str = Field("scamper", description="scamper, triz, six_thinking_hats, first_principles, biomimicry")
|
||||
num_ideas: int = Field(5, ge=1, le=20)
|
||||
constraints: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class SynthesisRequest(BaseModel):
|
||||
agent_id: str
|
||||
primary_domain: str
|
||||
secondary_domains: list[str]
|
||||
synthesis_goal: str
|
||||
|
||||
|
||||
# Endpoints
|
||||
|
||||
|
||||
@router.post("/capabilities", response_model=CreativeCapabilityResponse)
|
||||
async def create_creative_capability(request: CreativeCapabilityCreate, session: Annotated[Session, Depends(get_session)]) -> CreativeCapabilityResponse:
|
||||
"""Initialize a new creative capability for an agent"""
|
||||
engine = CreativityEnhancementEngine()
|
||||
|
||||
try:
|
||||
capability = await engine.create_creative_capability(
|
||||
session=session,
|
||||
agent_id=request.agent_id,
|
||||
creative_domain=request.creative_domain,
|
||||
capability_type=request.capability_type,
|
||||
generation_models=request.generation_models,
|
||||
initial_score=request.initial_score,
|
||||
)
|
||||
|
||||
return capability
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating creative capability: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/capabilities/{capability_id}/enhance")
|
||||
async def enhance_creativity(
|
||||
capability_id: str, request: EnhanceCreativityRequest, session: Annotated[Session, Depends(get_session)]
|
||||
) -> dict[str, Any]:
|
||||
"""Enhance a specific creative capability using specified algorithm"""
|
||||
engine = CreativityEnhancementEngine()
|
||||
|
||||
try:
|
||||
result = await engine.enhance_creativity(
|
||||
session=session, capability_id=capability_id, algorithm=request.algorithm, training_cycles=request.training_cycles
|
||||
)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error enhancing creativity: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/capabilities/{capability_id}/evaluate")
|
||||
async def evaluate_creation(
|
||||
capability_id: str, request: EvaluateCreationRequest, session: Annotated[Session, Depends(get_session)]
|
||||
) -> dict[str, Any]:
|
||||
"""Evaluate a creative output and update agent capability metrics"""
|
||||
engine = CreativityEnhancementEngine()
|
||||
|
||||
try:
|
||||
result = await engine.evaluate_creation(
|
||||
session=session,
|
||||
capability_id=capability_id,
|
||||
creation_data=request.creation_data,
|
||||
expert_feedback=request.expert_feedback,
|
||||
)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error evaluating creation: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/ideation/generate")
|
||||
async def generate_ideas(request: IdeationRequest) -> dict[str, Any]:
|
||||
"""Generate innovative ideas using specialized ideation algorithms"""
|
||||
ideation_engine = IdeationAlgorithm()
|
||||
|
||||
try:
|
||||
result = await ideation_engine.generate_ideas(
|
||||
problem_statement=request.problem_statement,
|
||||
domain=request.domain,
|
||||
technique=request.technique,
|
||||
num_ideas=request.num_ideas,
|
||||
constraints=request.constraints,
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating ideas: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/synthesis/cross-domain")
|
||||
async def synthesize_cross_domain(request: SynthesisRequest, session: Annotated[Session, Depends(get_session)]) -> dict[str, Any]:
|
||||
"""Synthesize concepts from multiple domains to create novel outputs"""
|
||||
integrator = CrossDomainCreativeIntegrator()
|
||||
|
||||
try:
|
||||
result = await integrator.generate_cross_domain_synthesis(
|
||||
session=session,
|
||||
agent_id=request.agent_id,
|
||||
primary_domain=request.primary_domain,
|
||||
secondary_domains=request.secondary_domains,
|
||||
synthesis_goal=request.synthesis_goal,
|
||||
)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error in cross-domain synthesis: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/capabilities/{agent_id}")
|
||||
async def list_agent_creative_capabilities(agent_id: str, session: Annotated[Session, Depends(get_session)]) -> list[CreativeCapability]:
|
||||
"""List all creative capabilities for a specific agent"""
|
||||
try:
|
||||
capabilities = session.execute(select(CreativeCapability).where(CreativeCapability.agent_id == agent_id)).all()
|
||||
|
||||
return capabilities
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching creative capabilities: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
570
apps/agent-management/src/app/routers/agent_integration_router.py
Executable file
570
apps/agent-management/src/app/routers/agent_integration_router.py
Executable file
@@ -0,0 +1,570 @@
|
||||
from typing import Annotated
|
||||
|
||||
"""
|
||||
Agent Integration and Deployment API Router for Verifiable AI Agent Orchestration
|
||||
Provides REST API endpoints for production deployment and integration management
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from ..deps import require_admin_key
|
||||
from app.domain.agent import AgentExecution, AIAgentWorkflow, VerificationLevel
|
||||
from ..services.agent_integration import (
|
||||
AgentDeploymentConfig,
|
||||
AgentDeploymentInstance,
|
||||
AgentDeploymentManager,
|
||||
AgentIntegrationManager,
|
||||
AgentMonitoringManager,
|
||||
AgentProductionManager,
|
||||
DeploymentStatus,
|
||||
)
|
||||
from ..storage import get_session
|
||||
from ..utils.alerting import alert_dispatcher
|
||||
|
||||
router = APIRouter(prefix="/agents/integration", tags=["Agent Integration"])
|
||||
|
||||
|
||||
@router.post("/deployments/config", response_model=AgentDeploymentConfig)
|
||||
async def create_deployment_config(
|
||||
workflow_id: str,
|
||||
deployment_name: str,
|
||||
deployment_config: dict,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentDeploymentConfig:
|
||||
"""Create deployment configuration for agent workflow"""
|
||||
|
||||
try:
|
||||
# Verify workflow exists and user has access
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
if workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
deployment_manager = AgentDeploymentManager(session)
|
||||
config = await deployment_manager.create_deployment_config(
|
||||
workflow_id=workflow_id, deployment_name=deployment_name, deployment_config=deployment_config
|
||||
)
|
||||
|
||||
logger.info("Deployment config created by %s", current_user)
|
||||
return config
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("Failed to create deployment config: %s", e)
|
||||
raise HTTPException(status_code=500, detail="Failed to create deployment config")
|
||||
|
||||
|
||||
@router.get("/deployments/configs", response_model=list[AgentDeploymentConfig])
|
||||
async def list_deployment_configs(
|
||||
workflow_id: str | None = None,
|
||||
status: DeploymentStatus | None = None,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> list[AgentDeploymentConfig]:
|
||||
"""List deployment configurations with filtering"""
|
||||
|
||||
try:
|
||||
query = select(AgentDeploymentConfig)
|
||||
|
||||
if workflow_id:
|
||||
query = query.where(AgentDeploymentConfig.workflow_id == workflow_id)
|
||||
|
||||
if status:
|
||||
query = query.where(AgentDeploymentConfig.status == status)
|
||||
|
||||
configs = session.execute(query).all()
|
||||
|
||||
# Filter by user ownership
|
||||
user_configs = []
|
||||
for config in configs:
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if workflow and workflow.owner_id == current_user:
|
||||
user_configs.append(config)
|
||||
|
||||
return user_configs
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list deployment configs: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/deployments/configs/{config_id}", response_model=AgentDeploymentConfig)
|
||||
async def get_deployment_config(
|
||||
config_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentDeploymentConfig:
|
||||
"""Get specific deployment configuration"""
|
||||
|
||||
try:
|
||||
config = session.get(AgentDeploymentConfig, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Deployment config not found")
|
||||
|
||||
# Check ownership
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
return config
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get deployment config: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/deployments/{config_id}/deploy")
|
||||
async def deploy_workflow(
|
||||
config_id: str,
|
||||
target_environment: str = "production",
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Deploy agent workflow to target environment"""
|
||||
|
||||
try:
|
||||
# Check ownership
|
||||
config = session.get(AgentDeploymentConfig, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Deployment config not found")
|
||||
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
deployment_manager = AgentDeploymentManager(session)
|
||||
deployment_result = await deployment_manager.deploy_agent_workflow(
|
||||
deployment_config_id=config_id, target_environment=target_environment
|
||||
)
|
||||
|
||||
logger.info(f"Workflow deployed: {config_id} to {target_environment} by {current_user}")
|
||||
return deployment_result
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to deploy workflow: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/deployments/{config_id}/health")
|
||||
async def get_deployment_health(
|
||||
config_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Get health status of deployment"""
|
||||
|
||||
try:
|
||||
# Check ownership
|
||||
config = session.get(AgentDeploymentConfig, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Deployment config not found")
|
||||
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
deployment_manager = AgentDeploymentManager(session)
|
||||
health_result = await deployment_manager.monitor_deployment_health(config_id)
|
||||
|
||||
return health_result
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get deployment health: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/deployments/{config_id}/scale")
|
||||
async def scale_deployment(
|
||||
config_id: str,
|
||||
target_instances: int,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Scale deployment to target number of instances"""
|
||||
|
||||
try:
|
||||
# Check ownership
|
||||
config = session.get(AgentDeploymentConfig, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Deployment config not found")
|
||||
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
deployment_manager = AgentDeploymentManager(session)
|
||||
scaling_result = await deployment_manager.scale_deployment(
|
||||
deployment_config_id=config_id, target_instances=target_instances
|
||||
)
|
||||
|
||||
logger.info(f"Deployment scaled: {config_id} to {target_instances} instances by {current_user}")
|
||||
return scaling_result
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to scale deployment: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/deployments/{config_id}/rollback")
|
||||
async def rollback_deployment(
|
||||
config_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Rollback deployment to previous version"""
|
||||
|
||||
try:
|
||||
# Check ownership
|
||||
config = session.get(AgentDeploymentConfig, config_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Deployment config not found")
|
||||
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
deployment_manager = AgentDeploymentManager(session)
|
||||
rollback_result = await deployment_manager.rollback_deployment(config_id)
|
||||
|
||||
logger.info(f"Deployment rolled back: {config_id} by {current_user}")
|
||||
return rollback_result
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to rollback deployment: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/deployments/instances", response_model=list[AgentDeploymentInstance])
|
||||
async def list_deployment_instances(
|
||||
deployment_id: str | None = None,
|
||||
environment: str | None = None,
|
||||
status: DeploymentStatus | None = None,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> list[AgentDeploymentInstance]:
|
||||
"""List deployment instances with filtering"""
|
||||
|
||||
try:
|
||||
query = select(AgentDeploymentInstance)
|
||||
|
||||
if deployment_id:
|
||||
query = query.where(AgentDeploymentInstance.deployment_id == deployment_id)
|
||||
|
||||
if environment:
|
||||
query = query.where(AgentDeploymentInstance.environment == environment)
|
||||
|
||||
if status:
|
||||
query = query.where(AgentDeploymentInstance.status == status)
|
||||
|
||||
instances = session.execute(query).all()
|
||||
|
||||
# Filter by user ownership
|
||||
user_instances = []
|
||||
for instance in instances:
|
||||
config = session.get(AgentDeploymentConfig, instance.deployment_id)
|
||||
if config:
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if workflow and workflow.owner_id == current_user:
|
||||
user_instances.append(instance)
|
||||
|
||||
return user_instances
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list deployment instances: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/deployments/instances/{instance_id}", response_model=AgentDeploymentInstance)
|
||||
async def get_deployment_instance(
|
||||
instance_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentDeploymentInstance:
|
||||
"""Get specific deployment instance"""
|
||||
|
||||
try:
|
||||
instance = session.get(AgentDeploymentInstance, instance_id)
|
||||
if not instance:
|
||||
raise HTTPException(status_code=404, detail="Instance not found")
|
||||
|
||||
# Check ownership
|
||||
config = session.get(AgentDeploymentConfig, instance.deployment_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Deployment config not found")
|
||||
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
return instance
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get deployment instance: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/integrations/zk/{execution_id}")
|
||||
async def integrate_with_zk_system(
|
||||
execution_id: str,
|
||||
verification_level: VerificationLevel = VerificationLevel.BASIC,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Integrate agent execution with ZK proof system"""
|
||||
|
||||
try:
|
||||
# Check execution ownership
|
||||
execution = session.get(AgentExecution, execution_id)
|
||||
if not execution:
|
||||
raise HTTPException(status_code=404, detail="Execution not found")
|
||||
|
||||
workflow = session.get(AIAgentWorkflow, execution.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
integration_manager = AgentIntegrationManager(session)
|
||||
integration_result = await integration_manager.integrate_with_zk_system(
|
||||
execution_id=execution_id, verification_level=verification_level
|
||||
)
|
||||
|
||||
logger.info(f"ZK integration completed: {execution_id} by {current_user}")
|
||||
return integration_result
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to integrate with ZK system: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/metrics/deployments/{deployment_id}")
|
||||
async def get_deployment_metrics(
|
||||
deployment_id: str,
|
||||
time_range: str = "1h",
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Get metrics for deployment over time range"""
|
||||
|
||||
try:
|
||||
# Check ownership
|
||||
config = session.get(AgentDeploymentConfig, deployment_id)
|
||||
if not config:
|
||||
raise HTTPException(status_code=404, detail="Deployment config not found")
|
||||
|
||||
workflow = session.get(AIAgentWorkflow, config.workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
monitoring_manager = AgentMonitoringManager(session)
|
||||
metrics = await monitoring_manager.get_deployment_metrics(deployment_config_id=deployment_id, time_range=time_range)
|
||||
|
||||
return metrics
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get deployment metrics: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/production/deploy")
|
||||
async def deploy_to_production(
|
||||
workflow_id: str,
|
||||
deployment_config: dict,
|
||||
integration_config: dict | None = None,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Deploy agent workflow to production with full integration"""
|
||||
|
||||
try:
|
||||
# Check workflow ownership
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
if workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
production_manager = AgentProductionManager(session)
|
||||
production_result = await production_manager.deploy_to_production(
|
||||
workflow_id=workflow_id, deployment_config=deployment_config, integration_config=integration_config
|
||||
)
|
||||
|
||||
logger.info(f"Production deployment completed: {workflow_id} by {current_user}")
|
||||
return production_result
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to deploy to production: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/production/dashboard")
|
||||
async def get_production_dashboard(
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]), current_user: str = Depends(require_admin_key())
|
||||
) -> dict[str, Any]:
|
||||
"""Get comprehensive production dashboard data"""
|
||||
|
||||
try:
|
||||
# Get user's deployments
|
||||
user_configs = session.execute(
|
||||
select(AgentDeploymentConfig).join(AIAgentWorkflow).where(AIAgentWorkflow.owner_id == current_user)
|
||||
).all()
|
||||
|
||||
dashboard_data = {
|
||||
"total_deployments": len(user_configs),
|
||||
"active_deployments": len([c for c in user_configs if c.status == DeploymentStatus.DEPLOYED]),
|
||||
"failed_deployments": len([c for c in user_configs if c.status == DeploymentStatus.FAILED]),
|
||||
"deployments": [],
|
||||
}
|
||||
|
||||
# Get detailed deployment info
|
||||
for config in user_configs:
|
||||
# Get instances for this deployment
|
||||
instances = session.execute(
|
||||
select(AgentDeploymentInstance).where(AgentDeploymentInstance.deployment_id == config.id)
|
||||
).all()
|
||||
|
||||
# Get metrics for this deployment
|
||||
try:
|
||||
monitoring_manager = AgentMonitoringManager(session)
|
||||
metrics = await monitoring_manager.get_deployment_metrics(config.id)
|
||||
except Exception:
|
||||
metrics = {"aggregated_metrics": {}}
|
||||
|
||||
dashboard_data["deployments"].append(
|
||||
{
|
||||
"deployment_id": config.id,
|
||||
"deployment_name": config.deployment_name,
|
||||
"workflow_id": config.workflow_id,
|
||||
"status": config.status,
|
||||
"total_instances": len(instances),
|
||||
"healthy_instances": len([i for i in instances if i.health_status == "healthy"]),
|
||||
"metrics": metrics["aggregated_metrics"],
|
||||
"created_at": config.created_at.isoformat(),
|
||||
"deployment_time": config.deployment_time.isoformat() if config.deployment_time else None,
|
||||
}
|
||||
)
|
||||
|
||||
return dashboard_data
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get production dashboard: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/production/health")
|
||||
async def get_production_health(
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]), current_user: str = Depends(require_admin_key())
|
||||
) -> dict[str, Any]:
|
||||
"""Get overall production health status"""
|
||||
|
||||
try:
|
||||
# Get user's deployments
|
||||
user_configs = session.execute(
|
||||
select(AgentDeploymentConfig).join(AIAgentWorkflow).where(AIAgentWorkflow.owner_id == current_user)
|
||||
).all()
|
||||
|
||||
health_status = {
|
||||
"overall_health": "healthy",
|
||||
"total_deployments": len(user_configs),
|
||||
"healthy_deployments": 0,
|
||||
"unhealthy_deployments": 0,
|
||||
"unknown_deployments": 0,
|
||||
"total_instances": 0,
|
||||
"healthy_instances": 0,
|
||||
"unhealthy_instances": 0,
|
||||
"deployment_health": [],
|
||||
}
|
||||
|
||||
# Check health of each deployment
|
||||
for config in user_configs:
|
||||
try:
|
||||
deployment_manager = AgentDeploymentManager(session)
|
||||
deployment_health = await deployment_manager.monitor_deployment_health(config.id)
|
||||
|
||||
health_status["deployment_health"].append(
|
||||
{
|
||||
"deployment_id": config.id,
|
||||
"deployment_name": config.deployment_name,
|
||||
"overall_health": deployment_health["overall_health"],
|
||||
"healthy_instances": deployment_health["healthy_instances"],
|
||||
"unhealthy_instances": deployment_health["unhealthy_instances"],
|
||||
"total_instances": deployment_health["total_instances"],
|
||||
}
|
||||
)
|
||||
|
||||
# Aggregate health counts
|
||||
health_status["total_instances"] += deployment_health["total_instances"]
|
||||
health_status["healthy_instances"] += deployment_health["healthy_instances"]
|
||||
health_status["unhealthy_instances"] += deployment_health["unhealthy_instances"]
|
||||
|
||||
if deployment_health["overall_health"] == "healthy":
|
||||
health_status["healthy_deployments"] += 1
|
||||
elif deployment_health["overall_health"] == "unhealthy":
|
||||
health_status["unhealthy_deployments"] += 1
|
||||
else:
|
||||
health_status["unknown_deployments"] += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Health check failed for deployment {config.id}: {e}")
|
||||
health_status["unknown_deployments"] += 1
|
||||
|
||||
# Determine overall health
|
||||
if health_status["unhealthy_deployments"] > 0:
|
||||
health_status["overall_health"] = "unhealthy"
|
||||
elif health_status["unknown_deployments"] > 0:
|
||||
health_status["overall_health"] = "degraded"
|
||||
|
||||
return health_status
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get production health: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/production/alerts")
|
||||
async def get_production_alerts(
|
||||
severity: str | None = None,
|
||||
limit: int = 50,
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Get production alerts and notifications"""
|
||||
|
||||
try:
|
||||
alerts = alert_dispatcher.get_recent_alerts(severity=severity, limit=limit)
|
||||
return {
|
||||
"alerts": alerts,
|
||||
"total_count": len(alerts),
|
||||
"severity": severity,
|
||||
"source": "coordinator_metrics",
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get production alerts: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
729
apps/agent-management/src/app/routers/agent_performance.py
Executable file
729
apps/agent-management/src/app/routers/agent_performance.py
Executable file
@@ -0,0 +1,729 @@
|
||||
from typing import Annotated
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
"""
|
||||
Advanced Agent Performance API Endpoints
|
||||
REST API for meta-learning, resource optimization, and performance enhancement
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from app.domain.agent_performance import (
|
||||
AgentCapability,
|
||||
AgentPerformanceProfile,
|
||||
CreativeCapability,
|
||||
FusionModel,
|
||||
LearningStrategy,
|
||||
MetaLearningModel,
|
||||
OptimizationTarget,
|
||||
PerformanceMetric,
|
||||
PerformanceOptimization,
|
||||
ReinforcementLearningConfig,
|
||||
ResourceAllocation,
|
||||
ResourceType,
|
||||
)
|
||||
from ..services.agent_performance_service import (
|
||||
AgentPerformanceService,
|
||||
MetaLearningEngine,
|
||||
PerformanceOptimizer,
|
||||
ResourceManager,
|
||||
)
|
||||
from ..storage import get_session
|
||||
|
||||
router = APIRouter(prefix="/v1/agent-performance", tags=["agent-performance"])
|
||||
|
||||
|
||||
# Pydantic models for API requests/responses
|
||||
class PerformanceProfileRequest(BaseModel):
|
||||
"""Request model for performance profile creation"""
|
||||
|
||||
agent_id: str
|
||||
agent_type: str = Field(default="hermes")
|
||||
initial_metrics: Dict[str, float] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class PerformanceProfileResponse(BaseModel):
|
||||
"""Response model for performance profile"""
|
||||
|
||||
profile_id: str
|
||||
agent_id: str
|
||||
agent_type: str
|
||||
overall_score: float
|
||||
performance_metrics: Dict[str, float]
|
||||
learning_strategies: List[str]
|
||||
specialization_areas: List[str]
|
||||
expertise_levels: Dict[str, float]
|
||||
resource_efficiency: Dict[str, float]
|
||||
cost_per_task: float
|
||||
throughput: float
|
||||
average_latency: float
|
||||
last_assessed: Optional[str]
|
||||
created_at: str
|
||||
updated_at: str
|
||||
|
||||
|
||||
class MetaLearningRequest(BaseModel):
|
||||
"""Request model for meta-learning model creation"""
|
||||
|
||||
model_name: str
|
||||
base_algorithms: List[str]
|
||||
meta_strategy: LearningStrategy
|
||||
adaptation_targets: List[str]
|
||||
|
||||
|
||||
class MetaLearningResponse(BaseModel):
|
||||
"""Response model for meta-learning model"""
|
||||
|
||||
model_id: str
|
||||
model_name: str
|
||||
model_type: str
|
||||
meta_strategy: str
|
||||
adaptation_targets: List[str]
|
||||
meta_accuracy: float
|
||||
adaptation_speed: float
|
||||
generalization_ability: float
|
||||
status: str
|
||||
created_at: str
|
||||
trained_at: Optional[str]
|
||||
|
||||
|
||||
class ResourceAllocationRequest(BaseModel):
|
||||
"""Request model for resource allocation"""
|
||||
|
||||
agent_id: str
|
||||
task_requirements: Dict[str, Any]
|
||||
optimization_target: OptimizationTarget = Field(default=OptimizationTarget.EFFICIENCY)
|
||||
priority_level: str = Field(default="normal")
|
||||
|
||||
|
||||
class ResourceAllocationResponse(BaseModel):
|
||||
"""Response model for resource allocation"""
|
||||
|
||||
allocation_id: str
|
||||
agent_id: str
|
||||
cpu_cores: float
|
||||
memory_gb: float
|
||||
gpu_count: float
|
||||
gpu_memory_gb: float
|
||||
storage_gb: float
|
||||
network_bandwidth: float
|
||||
optimization_target: str
|
||||
status: str
|
||||
allocated_at: str
|
||||
|
||||
|
||||
class PerformanceOptimizationRequest(BaseModel):
|
||||
"""Request model for performance optimization"""
|
||||
|
||||
agent_id: str
|
||||
target_metric: PerformanceMetric
|
||||
current_performance: Dict[str, float]
|
||||
optimization_type: str = Field(default="comprehensive")
|
||||
|
||||
|
||||
class PerformanceOptimizationResponse(BaseModel):
|
||||
"""Response model for performance optimization"""
|
||||
|
||||
optimization_id: str
|
||||
agent_id: str
|
||||
optimization_type: str
|
||||
target_metric: str
|
||||
status: str
|
||||
performance_improvement: float
|
||||
resource_savings: float
|
||||
cost_savings: float
|
||||
overall_efficiency_gain: float
|
||||
created_at: str
|
||||
completed_at: Optional[str]
|
||||
|
||||
|
||||
class CapabilityRequest(BaseModel):
|
||||
"""Request model for agent capability"""
|
||||
|
||||
agent_id: str
|
||||
capability_name: str
|
||||
capability_type: str
|
||||
domain_area: str
|
||||
skill_level: float = Field(ge=0, le=10.0)
|
||||
specialization_areas: List[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CapabilityResponse(BaseModel):
|
||||
"""Response model for agent capability"""
|
||||
|
||||
capability_id: str
|
||||
agent_id: str
|
||||
capability_name: str
|
||||
capability_type: str
|
||||
domain_area: str
|
||||
skill_level: float
|
||||
proficiency_score: float
|
||||
specialization_areas: List[str]
|
||||
status: str
|
||||
created_at: str
|
||||
|
||||
|
||||
# API Endpoints
|
||||
|
||||
|
||||
@router.post("/profiles", response_model=PerformanceProfileResponse)
|
||||
async def create_performance_profile(
|
||||
profile_request: PerformanceProfileRequest, session: Annotated[Session, Depends(get_session)]
|
||||
) -> PerformanceProfileResponse:
|
||||
"""Create agent performance profile"""
|
||||
|
||||
performance_service = AgentPerformanceService(session)
|
||||
|
||||
try:
|
||||
profile = await performance_service.create_performance_profile(
|
||||
agent_id=profile_request.agent_id,
|
||||
agent_type=profile_request.agent_type,
|
||||
initial_metrics=profile_request.initial_metrics,
|
||||
)
|
||||
|
||||
return PerformanceProfileResponse(
|
||||
profile_id=profile.profile_id,
|
||||
agent_id=profile.agent_id,
|
||||
agent_type=profile.agent_type,
|
||||
overall_score=profile.overall_score,
|
||||
performance_metrics=profile.performance_metrics,
|
||||
learning_strategies=profile.learning_strategies,
|
||||
specialization_areas=profile.specialization_areas,
|
||||
expertise_levels=profile.expertise_levels,
|
||||
resource_efficiency=profile.resource_efficiency,
|
||||
cost_per_task=profile.cost_per_task,
|
||||
throughput=profile.throughput,
|
||||
average_latency=profile.average_latency,
|
||||
last_assessed=profile.last_assessed.isoformat() if profile.last_assessed else None,
|
||||
created_at=profile.created_at.isoformat(),
|
||||
updated_at=profile.updated_at.isoformat(),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating performance profile: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.get("/profiles/{agent_id}", response_model=Dict[str, Any])
|
||||
async def get_performance_profile(agent_id: str, session: Annotated[Session, Depends(get_session)]) -> Dict[str, Any]:
|
||||
"""Get agent performance profile"""
|
||||
|
||||
performance_service = AgentPerformanceService(session)
|
||||
|
||||
try:
|
||||
profile = await performance_service.get_comprehensive_profile(agent_id)
|
||||
|
||||
if "error" in profile:
|
||||
raise HTTPException(status_code=404, detail=profile["error"])
|
||||
|
||||
return profile
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting performance profile for agent {agent_id}: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.post("/profiles/{agent_id}/metrics")
|
||||
async def update_performance_metrics(
|
||||
agent_id: str,
|
||||
metrics: Dict[str, float],
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
task_context: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Update agent performance metrics"""
|
||||
|
||||
performance_service = AgentPerformanceService(session)
|
||||
|
||||
try:
|
||||
profile = await performance_service.update_performance_metrics(
|
||||
agent_id=agent_id, new_metrics=metrics, task_context=task_context
|
||||
)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"profile_id": profile.profile_id,
|
||||
"overall_score": profile.overall_score,
|
||||
"updated_at": profile.updated_at.isoformat(),
|
||||
"improvement_trends": profile.improvement_trends,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating performance metrics for agent {agent_id}: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.post("/meta-learning/models", response_model=MetaLearningResponse)
|
||||
async def create_meta_learning_model(
|
||||
model_request: MetaLearningRequest, session: Annotated[Session, Depends(get_session)]
|
||||
) -> MetaLearningResponse:
|
||||
"""Create meta-learning model"""
|
||||
|
||||
meta_learning_engine = MetaLearningEngine()
|
||||
|
||||
try:
|
||||
model = await meta_learning_engine.create_meta_learning_model(
|
||||
session=session,
|
||||
model_name=model_request.model_name,
|
||||
base_algorithms=model_request.base_algorithms,
|
||||
meta_strategy=model_request.meta_strategy,
|
||||
adaptation_targets=model_request.adaptation_targets,
|
||||
)
|
||||
|
||||
return MetaLearningResponse(
|
||||
model_id=model.model_id,
|
||||
model_name=model.model_name,
|
||||
model_type=model.model_type,
|
||||
meta_strategy=model.meta_strategy.value,
|
||||
adaptation_targets=model.adaptation_targets,
|
||||
meta_accuracy=model.meta_accuracy,
|
||||
adaptation_speed=model.adaptation_speed,
|
||||
generalization_ability=model.generalization_ability,
|
||||
status=model.status,
|
||||
created_at=model.created_at.isoformat(),
|
||||
trained_at=model.trained_at.isoformat() if model.trained_at else None,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating meta-learning model: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.post("/meta-learning/models/{model_id}/adapt")
|
||||
async def adapt_model_to_task(
|
||||
model_id: str,
|
||||
task_data: Dict[str, Any],
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
adaptation_steps: int = Query(default=10, ge=1, le=50),
|
||||
) -> Dict[str, Any]:
|
||||
"""Adapt meta-learning model to new task"""
|
||||
|
||||
meta_learning_engine = MetaLearningEngine()
|
||||
|
||||
try:
|
||||
results = await meta_learning_engine.adapt_to_new_task(
|
||||
session=session, model_id=model_id, task_data=task_data, adaptation_steps=adaptation_steps
|
||||
)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"model_id": model_id,
|
||||
"adaptation_results": results,
|
||||
"adapted_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error adapting model {model_id}: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.get("/meta-learning/models")
|
||||
async def list_meta_learning_models(
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
status: Optional[str] = Query(default=None, description="Filter by status"),
|
||||
meta_strategy: Optional[str] = Query(default=None, description="Filter by meta strategy"),
|
||||
limit: int = Query(default=50, ge=1, le=100, description="Number of results"),
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""List meta-learning models"""
|
||||
|
||||
try:
|
||||
query = select(MetaLearningModel)
|
||||
|
||||
if status:
|
||||
query = query.where(MetaLearningModel.status == status)
|
||||
if meta_strategy:
|
||||
query = query.where(MetaLearningModel.meta_strategy == LearningStrategy(meta_strategy))
|
||||
|
||||
models = session.execute(query.order_by(MetaLearningModel.created_at.desc()).limit(limit)).all()
|
||||
|
||||
return [
|
||||
{
|
||||
"model_id": model.model_id,
|
||||
"model_name": model.model_name,
|
||||
"model_type": model.model_type,
|
||||
"meta_strategy": model.meta_strategy.value,
|
||||
"adaptation_targets": model.adaptation_targets,
|
||||
"meta_accuracy": model.meta_accuracy,
|
||||
"adaptation_speed": model.adaptation_speed,
|
||||
"generalization_ability": model.generalization_ability,
|
||||
"status": model.status,
|
||||
"deployment_count": model.deployment_count,
|
||||
"success_rate": model.success_rate,
|
||||
"created_at": model.created_at.isoformat(),
|
||||
"trained_at": model.trained_at.isoformat() if model.trained_at else None,
|
||||
}
|
||||
for model in models
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing meta-learning models: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.post("/resources/allocate", response_model=ResourceAllocationResponse)
|
||||
async def allocate_resources(
|
||||
allocation_request: ResourceAllocationRequest, session: Annotated[Session, Depends(get_session)]
|
||||
) -> ResourceAllocationResponse:
|
||||
"""Allocate resources for agent task"""
|
||||
|
||||
resource_manager = ResourceManager()
|
||||
|
||||
try:
|
||||
allocation = await resource_manager.allocate_resources(
|
||||
session=session,
|
||||
agent_id=allocation_request.agent_id,
|
||||
task_requirements=allocation_request.task_requirements,
|
||||
optimization_target=allocation_request.optimization_target,
|
||||
)
|
||||
|
||||
return ResourceAllocationResponse(
|
||||
allocation_id=allocation.allocation_id,
|
||||
agent_id=allocation.agent_id,
|
||||
cpu_cores=allocation.cpu_cores,
|
||||
memory_gb=allocation.memory_gb,
|
||||
gpu_count=allocation.gpu_count,
|
||||
gpu_memory_gb=allocation.gpu_memory_gb,
|
||||
storage_gb=allocation.storage_gb,
|
||||
network_bandwidth=allocation.network_bandwidth,
|
||||
optimization_target=allocation.optimization_target.value,
|
||||
status=allocation.status,
|
||||
allocated_at=allocation.allocated_at.isoformat(),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error allocating resources: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.get("/resources/{agent_id}")
|
||||
async def get_resource_allocations(
|
||||
agent_id: str,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
status: Optional[str] = Query(default=None, description="Filter by status"),
|
||||
limit: int = Query(default=20, ge=1, le=100, description="Number of results"),
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get resource allocations for agent"""
|
||||
|
||||
try:
|
||||
query = select(ResourceAllocation).where(ResourceAllocation.agent_id == agent_id)
|
||||
|
||||
if status:
|
||||
query = query.where(ResourceAllocation.status == status)
|
||||
|
||||
allocations = session.execute(query.order_by(ResourceAllocation.created_at.desc()).limit(limit)).all()
|
||||
|
||||
return [
|
||||
{
|
||||
"allocation_id": allocation.allocation_id,
|
||||
"agent_id": allocation.agent_id,
|
||||
"task_id": allocation.task_id,
|
||||
"cpu_cores": allocation.cpu_cores,
|
||||
"memory_gb": allocation.memory_gb,
|
||||
"gpu_count": allocation.gpu_count,
|
||||
"gpu_memory_gb": allocation.gpu_memory_gb,
|
||||
"storage_gb": allocation.storage_gb,
|
||||
"network_bandwidth": allocation.network_bandwidth,
|
||||
"optimization_target": allocation.optimization_target.value,
|
||||
"priority_level": allocation.priority_level,
|
||||
"status": allocation.status,
|
||||
"efficiency_score": allocation.efficiency_score,
|
||||
"cost_efficiency": allocation.cost_efficiency,
|
||||
"allocated_at": allocation.allocated_at.isoformat() if allocation.allocated_at else None,
|
||||
"started_at": allocation.started_at.isoformat() if allocation.started_at else None,
|
||||
"completed_at": allocation.completed_at.isoformat() if allocation.completed_at else None,
|
||||
}
|
||||
for allocation in allocations
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting resource allocations for agent {agent_id}: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.post("/optimization/optimize", response_model=PerformanceOptimizationResponse)
|
||||
async def optimize_performance(
|
||||
optimization_request: PerformanceOptimizationRequest, session: Annotated[Session, Depends(get_session)]
|
||||
) -> PerformanceOptimizationResponse:
|
||||
"""Optimize agent performance"""
|
||||
|
||||
performance_optimizer = PerformanceOptimizer()
|
||||
|
||||
try:
|
||||
optimization = await performance_optimizer.optimize_agent_performance(
|
||||
session=session,
|
||||
agent_id=optimization_request.agent_id,
|
||||
target_metric=optimization_request.target_metric,
|
||||
current_performance=optimization_request.current_performance,
|
||||
)
|
||||
|
||||
return PerformanceOptimizationResponse(
|
||||
optimization_id=optimization.optimization_id,
|
||||
agent_id=optimization.agent_id,
|
||||
optimization_type=optimization.optimization_type,
|
||||
target_metric=optimization.target_metric.value,
|
||||
status=optimization.status,
|
||||
performance_improvement=optimization.performance_improvement,
|
||||
resource_savings=optimization.resource_savings,
|
||||
cost_savings=optimization.cost_savings,
|
||||
overall_efficiency_gain=optimization.overall_efficiency_gain,
|
||||
created_at=optimization.created_at.isoformat(),
|
||||
completed_at=optimization.completed_at.isoformat() if optimization.completed_at else None,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error optimizing performance: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.get("/optimization/{agent_id}")
|
||||
async def get_optimization_history(
|
||||
agent_id: str,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
status: Optional[str] = Query(default=None, description="Filter by status"),
|
||||
target_metric: Optional[str] = Query(default=None, description="Filter by target metric"),
|
||||
limit: int = Query(default=20, ge=1, le=100, description="Number of results"),
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get optimization history for agent"""
|
||||
|
||||
try:
|
||||
query = select(PerformanceOptimization).where(PerformanceOptimization.agent_id == agent_id)
|
||||
|
||||
if status:
|
||||
query = query.where(PerformanceOptimization.status == status)
|
||||
if target_metric:
|
||||
query = query.where(PerformanceOptimization.target_metric == PerformanceMetric(target_metric))
|
||||
|
||||
optimizations = session.execute(query.order_by(PerformanceOptimization.created_at.desc()).limit(limit)).all()
|
||||
|
||||
return [
|
||||
{
|
||||
"optimization_id": optimization.optimization_id,
|
||||
"agent_id": optimization.agent_id,
|
||||
"optimization_type": optimization.optimization_type,
|
||||
"target_metric": optimization.target_metric.value,
|
||||
"status": optimization.status,
|
||||
"baseline_performance": optimization.baseline_performance,
|
||||
"optimized_performance": optimization.optimized_performance,
|
||||
"baseline_cost": optimization.baseline_cost,
|
||||
"optimized_cost": optimization.optimized_cost,
|
||||
"performance_improvement": optimization.performance_improvement,
|
||||
"resource_savings": optimization.resource_savings,
|
||||
"cost_savings": optimization.cost_savings,
|
||||
"overall_efficiency_gain": optimization.overall_efficiency_gain,
|
||||
"optimization_duration": optimization.optimization_duration,
|
||||
"iterations_required": optimization.iterations_required,
|
||||
"convergence_achieved": optimization.convergence_achieved,
|
||||
"created_at": optimization.created_at.isoformat(),
|
||||
"completed_at": optimization.completed_at.isoformat() if optimization.completed_at else None,
|
||||
}
|
||||
for optimization in optimizations
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting optimization history for agent {agent_id}: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.post("/capabilities", response_model=CapabilityResponse)
|
||||
async def create_capability(
|
||||
capability_request: CapabilityRequest, session: Annotated[Session, Depends(get_session)]
|
||||
) -> CapabilityResponse:
|
||||
"""Create agent capability"""
|
||||
|
||||
try:
|
||||
capability_id = f"cap_{uuid4().hex[:8]}"
|
||||
|
||||
capability = AgentCapability(
|
||||
capability_id=capability_id,
|
||||
agent_id=capability_request.agent_id,
|
||||
capability_name=capability_request.capability_name,
|
||||
capability_type=capability_request.capability_type,
|
||||
domain_area=capability_request.domain_area,
|
||||
skill_level=capability_request.skill_level,
|
||||
specialization_areas=capability_request.specialization_areas,
|
||||
proficiency_score=min(1.0, capability_request.skill_level / 10.0),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
session.add(capability)
|
||||
session.commit()
|
||||
session.refresh(capability)
|
||||
|
||||
return CapabilityResponse(
|
||||
capability_id=capability.capability_id,
|
||||
agent_id=capability.agent_id,
|
||||
capability_name=capability.capability_name,
|
||||
capability_type=capability.capability_type,
|
||||
domain_area=capability.domain_area,
|
||||
skill_level=capability.skill_level,
|
||||
proficiency_score=capability.proficiency_score,
|
||||
specialization_areas=capability.specialization_areas,
|
||||
status=capability.status,
|
||||
created_at=capability.created_at.isoformat(),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating capability: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.get("/capabilities/{agent_id}")
|
||||
async def get_agent_capabilities(
|
||||
agent_id: str,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
capability_type: Optional[str] = Query(default=None, description="Filter by capability type"),
|
||||
domain_area: Optional[str] = Query(default=None, description="Filter by domain area"),
|
||||
limit: int = Query(default=50, ge=1, le=100, description="Number of results"),
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get agent capabilities"""
|
||||
|
||||
try:
|
||||
query = select(AgentCapability).where(AgentCapability.agent_id == agent_id)
|
||||
|
||||
if capability_type:
|
||||
query = query.where(AgentCapability.capability_type == capability_type)
|
||||
if domain_area:
|
||||
query = query.where(AgentCapability.domain_area == domain_area)
|
||||
|
||||
capabilities = session.execute(query.order_by(AgentCapability.skill_level.desc()).limit(limit)).all()
|
||||
|
||||
return [
|
||||
{
|
||||
"capability_id": capability.capability_id,
|
||||
"agent_id": capability.agent_id,
|
||||
"capability_name": capability.capability_name,
|
||||
"capability_type": capability.capability_type,
|
||||
"domain_area": capability.domain_area,
|
||||
"skill_level": capability.skill_level,
|
||||
"proficiency_score": capability.proficiency_score,
|
||||
"experience_years": capability.experience_years,
|
||||
"success_rate": capability.success_rate,
|
||||
"average_quality": capability.average_quality,
|
||||
"learning_rate": capability.learning_rate,
|
||||
"adaptation_speed": capability.adaptation_speed,
|
||||
"specialization_areas": capability.specialization_areas,
|
||||
"sub_capabilities": capability.sub_capabilities,
|
||||
"tool_proficiency": capability.tool_proficiency,
|
||||
"certified": capability.certified,
|
||||
"certification_level": capability.certification_level,
|
||||
"status": capability.status,
|
||||
"acquired_at": capability.acquired_at.isoformat(),
|
||||
"last_improved": capability.last_improved.isoformat() if capability.last_improved else None,
|
||||
}
|
||||
for capability in capabilities
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting capabilities for agent {agent_id}: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
@router.get("/analytics/performance-summary")
|
||||
async def get_performance_summary(
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
agent_ids: List[str] = Query(default=[], description="List of agent IDs"),
|
||||
metric: Optional[str] = Query(default="overall_score", description="Metric to summarize"),
|
||||
period: str = Query(default="7d", description="Time period"),
|
||||
) -> Dict[str, Any]:
|
||||
"""Get performance summary for agents"""
|
||||
|
||||
try:
|
||||
if not agent_ids:
|
||||
# Get all agents if none specified
|
||||
profiles = session.execute(select(AgentPerformanceProfile)).all()
|
||||
agent_ids = [p.agent_id for p in profiles]
|
||||
|
||||
summaries = []
|
||||
|
||||
for agent_id in agent_ids:
|
||||
profile = session.execute(
|
||||
select(AgentPerformanceProfile).where(AgentPerformanceProfile.agent_id == agent_id)
|
||||
).first()
|
||||
|
||||
if profile:
|
||||
summaries.append(
|
||||
{
|
||||
"agent_id": agent_id,
|
||||
"overall_score": profile.overall_score,
|
||||
"performance_metrics": profile.performance_metrics,
|
||||
"resource_efficiency": profile.resource_efficiency,
|
||||
"cost_per_task": profile.cost_per_task,
|
||||
"throughput": profile.throughput,
|
||||
"average_latency": profile.average_latency,
|
||||
"specialization_areas": profile.specialization_areas,
|
||||
"last_assessed": profile.last_assessed.isoformat() if profile.last_assessed else None,
|
||||
}
|
||||
)
|
||||
|
||||
# Calculate summary statistics
|
||||
if summaries:
|
||||
overall_scores = [s["overall_score"] for s in summaries]
|
||||
avg_score = sum(overall_scores) / len(overall_scores)
|
||||
|
||||
return {
|
||||
"period": period,
|
||||
"agent_count": len(summaries),
|
||||
"average_score": avg_score,
|
||||
"top_performers": sorted(summaries, key=lambda x: x["overall_score"], reverse=True)[:10],
|
||||
"performance_distribution": {
|
||||
"excellent": len([s for s in summaries if s["overall_score"] >= 80]),
|
||||
"good": len([s for s in summaries if 60 <= s["overall_score"] < 80]),
|
||||
"average": len([s for s in summaries if 40 <= s["overall_score"] < 60]),
|
||||
"below_average": len([s for s in summaries if s["overall_score"] < 40]),
|
||||
},
|
||||
"specialization_distribution": self.calculate_specialization_distribution(summaries),
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"period": period,
|
||||
"agent_count": 0,
|
||||
"average_score": 0.0,
|
||||
"top_performers": [],
|
||||
"performance_distribution": {},
|
||||
"specialization_distribution": {},
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting performance summary: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
|
||||
def calculate_specialization_distribution(summaries: List[Dict[str, Any]]) -> Dict[str, int]:
|
||||
"""Calculate specialization distribution"""
|
||||
|
||||
distribution = {}
|
||||
|
||||
for summary in summaries:
|
||||
for area in summary["specialization_areas"]:
|
||||
distribution[area] = distribution.get(area, 0) + 1
|
||||
|
||||
return distribution
|
||||
|
||||
|
||||
@router.get("/health")
|
||||
async def health_check() -> Dict[str, Any]:
|
||||
"""Health check for agent performance service"""
|
||||
|
||||
return {
|
||||
"status": "healthy",
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"version": "1.0.0",
|
||||
"services": {
|
||||
"meta_learning_engine": "operational",
|
||||
"resource_manager": "operational",
|
||||
"performance_optimizer": "operational",
|
||||
"performance_service": "operational",
|
||||
},
|
||||
}
|
||||
506
apps/agent-management/src/app/routers/agent_router.py
Executable file
506
apps/agent-management/src/app/routers/agent_router.py
Executable file
@@ -0,0 +1,506 @@
|
||||
from typing import Annotated
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
"""
|
||||
AI Agent API Router for Verifiable AI Agent Orchestration
|
||||
Provides REST API endpoints for agent workflow management and execution
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from ..deps import require_admin_key
|
||||
from app.domain.agent import (
|
||||
AgentExecutionRequest,
|
||||
AgentExecutionResponse,
|
||||
AgentExecutionStatus,
|
||||
AgentStatus,
|
||||
AgentWorkflowCreate,
|
||||
AgentWorkflowUpdate,
|
||||
AIAgentWorkflow,
|
||||
)
|
||||
from ..services.agent_service import AIAgentOrchestrator
|
||||
from ..storage import get_session
|
||||
|
||||
router = APIRouter(tags=["AI Agents"])
|
||||
|
||||
|
||||
@router.post("/workflows", response_model=AIAgentWorkflow)
|
||||
async def create_workflow(
|
||||
workflow_data: AgentWorkflowCreate,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AIAgentWorkflow:
|
||||
"""Create a new AI agent workflow"""
|
||||
|
||||
try:
|
||||
workflow = AIAgentWorkflow(owner_id=current_user, **workflow_data.dict()) # Use string directly
|
||||
|
||||
session.add(workflow)
|
||||
session.commit()
|
||||
session.refresh(workflow)
|
||||
|
||||
logger.info(f"Created agent workflow: {workflow.id}")
|
||||
return workflow
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create workflow: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/workflows", response_model=list[AIAgentWorkflow])
|
||||
async def list_workflows(
|
||||
owner_id: str | None = None,
|
||||
is_public: bool | None = None,
|
||||
tags: list[str] | None = None,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> list[AIAgentWorkflow]:
|
||||
"""List agent workflows with filtering"""
|
||||
|
||||
try:
|
||||
query = select(AIAgentWorkflow)
|
||||
|
||||
# Filter by owner or public workflows
|
||||
if owner_id:
|
||||
query = query.where(AIAgentWorkflow.owner_id == owner_id)
|
||||
elif not is_public:
|
||||
query = query.where((AIAgentWorkflow.owner_id == current_user.id) | (AIAgentWorkflow.is_public))
|
||||
|
||||
# Filter by public status
|
||||
if is_public is not None:
|
||||
query = query.where(AIAgentWorkflow.is_public == is_public)
|
||||
|
||||
# Filter by tags
|
||||
if tags:
|
||||
for tag in tags:
|
||||
query = query.where(AIAgentWorkflow.tags.contains([tag]))
|
||||
|
||||
workflows = session.execute(query).all()
|
||||
return workflows
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list workflows: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/workflows/{workflow_id}", response_model=AIAgentWorkflow)
|
||||
async def get_workflow(
|
||||
workflow_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AIAgentWorkflow:
|
||||
"""Get a specific agent workflow"""
|
||||
|
||||
try:
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
# Check access permissions
|
||||
if workflow.owner_id != current_user and not workflow.is_public:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
return workflow
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get workflow: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.put("/workflows/{workflow_id}", response_model=AIAgentWorkflow)
|
||||
async def update_workflow(
|
||||
workflow_id: str,
|
||||
workflow_data: AgentWorkflowUpdate,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AIAgentWorkflow:
|
||||
"""Update an agent workflow"""
|
||||
|
||||
try:
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
# Check ownership
|
||||
if workflow.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
# Update workflow
|
||||
update_data = workflow_data.dict(exclude_unset=True)
|
||||
for field, value in update_data.items():
|
||||
setattr(workflow, field, value)
|
||||
|
||||
workflow.updated_at = datetime.now(timezone.utc)
|
||||
session.commit()
|
||||
session.refresh(workflow)
|
||||
|
||||
logger.info(f"Updated agent workflow: {workflow.id}")
|
||||
return workflow
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update workflow: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.delete("/workflows/{workflow_id}")
|
||||
async def delete_workflow(
|
||||
workflow_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, str]:
|
||||
"""Delete an agent workflow"""
|
||||
|
||||
try:
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
# Check ownership
|
||||
if workflow.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
session.delete(workflow)
|
||||
session.commit()
|
||||
|
||||
logger.info(f"Deleted agent workflow: {workflow_id}")
|
||||
return {"message": "Workflow deleted successfully"}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete workflow: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/workflows/{workflow_id}/execute", response_model=AgentExecutionResponse)
|
||||
async def execute_workflow(
|
||||
workflow_id: str,
|
||||
execution_request: AgentExecutionRequest,
|
||||
background_tasks: BackgroundTasks,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentExecutionResponse:
|
||||
"""Execute an AI agent workflow"""
|
||||
|
||||
try:
|
||||
# Verify workflow exists and user has access
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
if workflow.owner_id != current_user.id and not workflow.is_public:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
# Create execution request
|
||||
request = AgentExecutionRequest(
|
||||
workflow_id=workflow_id,
|
||||
inputs=execution_request.inputs,
|
||||
verification_level=execution_request.verification_level or workflow.verification_level,
|
||||
max_execution_time=execution_request.max_execution_time or workflow.max_execution_time,
|
||||
max_cost_budget=execution_request.max_cost_budget or workflow.max_cost_budget,
|
||||
)
|
||||
|
||||
# Create orchestrator and execute
|
||||
from ..coordinator_client import CoordinatorClient
|
||||
|
||||
coordinator_client = CoordinatorClient()
|
||||
orchestrator = AIAgentOrchestrator(session, coordinator_client)
|
||||
|
||||
response = await orchestrator.execute_workflow(request, current_user.id)
|
||||
|
||||
logger.info(f"Started agent execution: {response.execution_id}")
|
||||
return response
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to execute workflow: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/executions/{execution_id}/status", response_model=AgentExecutionStatus)
|
||||
async def get_execution_status(
|
||||
execution_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentExecutionStatus:
|
||||
"""Get execution status"""
|
||||
|
||||
try:
|
||||
from ..coordinator_client import CoordinatorClient
|
||||
from ..services.agent_service import AIAgentOrchestrator
|
||||
|
||||
coordinator_client = CoordinatorClient()
|
||||
orchestrator = AIAgentOrchestrator(session, coordinator_client)
|
||||
|
||||
status = await orchestrator.get_execution_status(execution_id)
|
||||
|
||||
# Verify user has access to this execution
|
||||
workflow = session.get(AIAgentWorkflow, status.workflow_id)
|
||||
if workflow.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
return status
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get execution status: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/executions", response_model=list[AgentExecutionStatus])
|
||||
async def list_executions(
|
||||
workflow_id: str | None = None,
|
||||
status: AgentStatus | None = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> list[AgentExecutionStatus]:
|
||||
"""List agent executions with filtering"""
|
||||
|
||||
try:
|
||||
from app.domain.agent import AgentExecution
|
||||
|
||||
query = select(AgentExecution)
|
||||
|
||||
# Filter by user's workflows
|
||||
if workflow_id:
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow or workflow.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
query = query.where(AgentExecution.workflow_id == workflow_id)
|
||||
else:
|
||||
# Get all workflows owned by user
|
||||
user_workflows = session.execute(
|
||||
select(AIAgentWorkflow.id).where(AIAgentWorkflow.owner_id == current_user.id)
|
||||
).all()
|
||||
workflow_ids = [w.id for w in user_workflows]
|
||||
query = query.where(AgentExecution.workflow_id.in_(workflow_ids))
|
||||
|
||||
# Filter by status
|
||||
if status:
|
||||
query = query.where(AgentExecution.status == status)
|
||||
|
||||
# Apply pagination
|
||||
query = query.offset(offset).limit(limit)
|
||||
query = query.order_by(AgentExecution.created_at.desc())
|
||||
|
||||
executions = session.execute(query).all()
|
||||
|
||||
# Convert to response models
|
||||
execution_statuses = []
|
||||
for execution in executions:
|
||||
from ..coordinator_client import CoordinatorClient
|
||||
from ..services.agent_service import AIAgentOrchestrator
|
||||
|
||||
coordinator_client = CoordinatorClient()
|
||||
orchestrator = AIAgentOrchestrator(session, coordinator_client)
|
||||
|
||||
status = await orchestrator.get_execution_status(execution.id)
|
||||
execution_statuses.append(status)
|
||||
|
||||
return execution_statuses
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list executions: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/executions/{execution_id}/cancel")
|
||||
async def cancel_execution(
|
||||
execution_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, str]:
|
||||
"""Cancel an ongoing execution"""
|
||||
|
||||
try:
|
||||
from app.domain.agent import AgentExecution
|
||||
from ..services.agent_service import AgentStateManager
|
||||
|
||||
# Get execution
|
||||
execution = session.get(AgentExecution, execution_id)
|
||||
if not execution:
|
||||
raise HTTPException(status_code=404, detail="Execution not found")
|
||||
|
||||
# Verify user has access
|
||||
workflow = session.get(AIAgentWorkflow, execution.workflow_id)
|
||||
if workflow.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
# Check if execution can be cancelled
|
||||
if execution.status not in [AgentStatus.PENDING, AgentStatus.RUNNING]:
|
||||
raise HTTPException(status_code=400, detail="Execution cannot be cancelled")
|
||||
|
||||
# Cancel execution
|
||||
state_manager = AgentStateManager(session)
|
||||
await state_manager.update_execution_status(execution_id, status=AgentStatus.CANCELLED, completed_at=datetime.now(timezone.utc))
|
||||
|
||||
logger.info(f"Cancelled agent execution: {execution_id}")
|
||||
return {"message": "Execution cancelled successfully"}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to cancel execution: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/executions/{execution_id}/logs")
|
||||
async def get_execution_logs(
|
||||
execution_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Get execution logs"""
|
||||
|
||||
try:
|
||||
from app.domain.agent import AgentExecution, AgentStepExecution
|
||||
|
||||
# Get execution
|
||||
execution = session.get(AgentExecution, execution_id)
|
||||
if not execution:
|
||||
raise HTTPException(status_code=404, detail="Execution not found")
|
||||
|
||||
# Verify user has access
|
||||
workflow = session.get(AIAgentWorkflow, execution.workflow_id)
|
||||
if workflow.owner_id != current_user.id:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
# Get step executions
|
||||
step_executions = session.execute(
|
||||
select(AgentStepExecution).where(AgentStepExecution.execution_id == execution_id)
|
||||
).all()
|
||||
|
||||
logs = []
|
||||
for step_exec in step_executions:
|
||||
logs.append(
|
||||
{
|
||||
"step_id": step_exec.step_id,
|
||||
"status": step_exec.status,
|
||||
"started_at": step_exec.started_at,
|
||||
"completed_at": step_exec.completed_at,
|
||||
"execution_time": step_exec.execution_time,
|
||||
"error_message": step_exec.error_message,
|
||||
"gpu_accelerated": step_exec.gpu_accelerated,
|
||||
"memory_usage": step_exec.memory_usage,
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"execution_id": execution_id,
|
||||
"workflow_id": execution.workflow_id,
|
||||
"status": execution.status,
|
||||
"started_at": execution.started_at,
|
||||
"completed_at": execution.completed_at,
|
||||
"total_execution_time": execution.total_execution_time,
|
||||
"step_logs": logs,
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get execution logs: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/test")
|
||||
async def test_endpoint() -> dict[str, str]:
|
||||
"""Test endpoint to verify router is working"""
|
||||
return {"message": "Agent router is working", "timestamp": datetime.now(timezone.utc).isoformat()}
|
||||
|
||||
|
||||
@router.post("/networks", response_model=dict, status_code=201)
|
||||
async def create_agent_network(
|
||||
network_data: dict,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Create a new agent network for collaborative processing"""
|
||||
|
||||
try:
|
||||
# Validate required fields
|
||||
if not network_data.get("name"):
|
||||
raise HTTPException(status_code=400, detail="Network name is required")
|
||||
|
||||
if not network_data.get("agents"):
|
||||
raise HTTPException(status_code=400, detail="Agent list is required")
|
||||
|
||||
# Create network record (simplified for now)
|
||||
network_id = f"network_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}"
|
||||
|
||||
network_response = {
|
||||
"id": network_id,
|
||||
"name": network_data["name"],
|
||||
"description": network_data.get("description", ""),
|
||||
"agents": network_data["agents"],
|
||||
"coordination_strategy": network_data.get("coordination", "centralized"),
|
||||
"status": "active",
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
"owner_id": current_user,
|
||||
}
|
||||
|
||||
logger.info(f"Created agent network: {network_id}")
|
||||
return network_response
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create agent network: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/executions/{execution_id}/receipt")
|
||||
async def get_execution_receipt(
|
||||
execution_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Get verifiable receipt for completed execution"""
|
||||
|
||||
try:
|
||||
# For now, return a mock receipt since the full execution system isn't implemented
|
||||
receipt_data = {
|
||||
"execution_id": execution_id,
|
||||
"workflow_id": f"workflow_{execution_id}",
|
||||
"status": "completed",
|
||||
"receipt_id": f"receipt_{execution_id}",
|
||||
"miner_signature": "0xmock_signature_placeholder",
|
||||
"coordinator_attestations": [
|
||||
{
|
||||
"coordinator_id": "coordinator_1",
|
||||
"signature": "0xmock_attestation_1",
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
],
|
||||
"minted_amount": 1000,
|
||||
"recorded_at": datetime.now(timezone.utc).isoformat(),
|
||||
"verified": True,
|
||||
"block_hash": "0xmock_block_hash",
|
||||
"transaction_hash": "0xmock_tx_hash",
|
||||
}
|
||||
|
||||
logger.info(f"Generated receipt for execution: {execution_id}")
|
||||
return receipt_data
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get execution receipt: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
650
apps/agent-management/src/app/routers/agent_security_router.py
Executable file
650
apps/agent-management/src/app/routers/agent_security_router.py
Executable file
@@ -0,0 +1,650 @@
|
||||
from typing import Annotated
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
"""
|
||||
Agent Security API Router for Verifiable AI Agent Orchestration
|
||||
Provides REST API endpoints for security management and auditing
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from ..deps import require_admin_key
|
||||
from app.domain.agent import AIAgentWorkflow
|
||||
from ..services.agent_security import (
|
||||
AgentAuditLog,
|
||||
AgentAuditor,
|
||||
AgentSandboxManager,
|
||||
AgentSecurityManager,
|
||||
AgentSecurityPolicy,
|
||||
AgentTrustManager,
|
||||
AgentTrustScore,
|
||||
AuditEventType,
|
||||
SecurityLevel,
|
||||
)
|
||||
from ..storage import get_session
|
||||
|
||||
router = APIRouter(prefix="/agents/security", tags=["Agent Security"])
|
||||
|
||||
|
||||
@router.post("/policies", response_model=AgentSecurityPolicy)
|
||||
async def create_security_policy(
|
||||
name: str,
|
||||
description: str,
|
||||
security_level: SecurityLevel,
|
||||
policy_rules: dict,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentSecurityPolicy:
|
||||
"""Create a new security policy"""
|
||||
|
||||
try:
|
||||
security_manager = AgentSecurityManager(session)
|
||||
policy = await security_manager.create_security_policy(
|
||||
name=name, description=description, security_level=security_level, policy_rules=policy_rules
|
||||
)
|
||||
|
||||
logger.info(f"Security policy created: {policy.id} by {current_user}")
|
||||
return policy
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create security policy: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/policies", response_model=list[AgentSecurityPolicy])
|
||||
async def list_security_policies(
|
||||
security_level: SecurityLevel | None = None,
|
||||
is_active: bool | None = None,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> list[AgentSecurityPolicy]:
|
||||
"""List security policies with filtering"""
|
||||
|
||||
try:
|
||||
query = select(AgentSecurityPolicy)
|
||||
|
||||
if security_level:
|
||||
query = query.where(AgentSecurityPolicy.security_level == security_level)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.where(AgentSecurityPolicy.is_active == is_active)
|
||||
|
||||
policies = session.execute(query).all()
|
||||
return policies
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list security policies: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/policies/{policy_id}", response_model=AgentSecurityPolicy)
|
||||
async def get_security_policy(
|
||||
policy_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentSecurityPolicy:
|
||||
"""Get a specific security policy"""
|
||||
|
||||
try:
|
||||
policy = session.get(AgentSecurityPolicy, policy_id)
|
||||
if not policy:
|
||||
raise HTTPException(status_code=404, detail="Policy not found")
|
||||
|
||||
return policy
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get security policy: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.put("/policies/{policy_id}", response_model=AgentSecurityPolicy)
|
||||
async def update_security_policy(
|
||||
policy_id: str,
|
||||
policy_updates: dict,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentSecurityPolicy:
|
||||
"""Update a security policy"""
|
||||
|
||||
try:
|
||||
policy = session.get(AgentSecurityPolicy, policy_id)
|
||||
if not policy:
|
||||
raise HTTPException(status_code=404, detail="Policy not found")
|
||||
|
||||
# Update policy fields
|
||||
for field, value in policy_updates.items():
|
||||
if hasattr(policy, field):
|
||||
setattr(policy, field, value)
|
||||
|
||||
policy.updated_at = datetime.now(timezone.utc)
|
||||
session.commit()
|
||||
session.refresh(policy)
|
||||
|
||||
# Log policy update
|
||||
auditor = AgentAuditor(session)
|
||||
await auditor.log_event(
|
||||
AuditEventType.WORKFLOW_UPDATED,
|
||||
user_id=current_user,
|
||||
security_level=policy.security_level,
|
||||
event_data={"policy_id": policy_id, "updates": policy_updates},
|
||||
new_state={"policy": policy.dict()},
|
||||
)
|
||||
|
||||
logger.info(f"Security policy updated: {policy_id} by {current_user}")
|
||||
return policy
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update security policy: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.delete("/policies/{policy_id}")
|
||||
async def delete_security_policy(
|
||||
policy_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, str]:
|
||||
"""Delete a security policy"""
|
||||
|
||||
try:
|
||||
policy = session.get(AgentSecurityPolicy, policy_id)
|
||||
if not policy:
|
||||
raise HTTPException(status_code=404, detail="Policy not found")
|
||||
|
||||
# Log policy deletion
|
||||
auditor = AgentAuditor(session)
|
||||
await auditor.log_event(
|
||||
AuditEventType.WORKFLOW_DELETED,
|
||||
user_id=current_user,
|
||||
security_level=policy.security_level,
|
||||
event_data={"policy_id": policy_id, "policy_name": policy.name},
|
||||
previous_state={"policy": policy.dict()},
|
||||
)
|
||||
|
||||
session.delete(policy)
|
||||
session.commit()
|
||||
|
||||
logger.info(f"Security policy deleted: {policy_id} by {current_user}")
|
||||
return {"message": "Policy deleted successfully"}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete security policy: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/validate-workflow/{workflow_id}")
|
||||
async def validate_workflow_security(
|
||||
workflow_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Validate workflow security requirements"""
|
||||
|
||||
try:
|
||||
workflow = session.get(AIAgentWorkflow, workflow_id)
|
||||
if not workflow:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
# Check ownership
|
||||
if workflow.owner_id != current_user:
|
||||
raise HTTPException(status_code=403, detail="Access denied")
|
||||
|
||||
security_manager = AgentSecurityManager(session)
|
||||
validation_result = await security_manager.validate_workflow_security(workflow, current_user)
|
||||
|
||||
return validation_result
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to validate workflow security: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/audit-logs", response_model=list[AgentAuditLog])
|
||||
async def list_audit_logs(
|
||||
event_type: AuditEventType | None = None,
|
||||
workflow_id: str | None = None,
|
||||
execution_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
security_level: SecurityLevel | None = None,
|
||||
requires_investigation: bool | None = None,
|
||||
risk_score_min: int | None = None,
|
||||
risk_score_max: int | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> list[AgentAuditLog]:
|
||||
"""List audit logs with filtering"""
|
||||
|
||||
try:
|
||||
from ..services.agent_security import AgentAuditLog
|
||||
|
||||
query = select(AgentAuditLog)
|
||||
|
||||
# Apply filters
|
||||
if event_type:
|
||||
query = query.where(AgentAuditLog.event_type == event_type)
|
||||
if workflow_id:
|
||||
query = query.where(AgentAuditLog.workflow_id == workflow_id)
|
||||
if execution_id:
|
||||
query = query.where(AgentLog.execution_id == execution_id)
|
||||
if user_id:
|
||||
query = query.where(AuditLog.user_id == user_id)
|
||||
if security_level:
|
||||
query = query.where(AuditLog.security_level == security_level)
|
||||
if requires_investigation is not None:
|
||||
query = query.where(AuditLog.requires_investigation == requires_investigation)
|
||||
if risk_score_min is not None:
|
||||
query = query.where(AuditLog.risk_score >= risk_score_min)
|
||||
if risk_score_max is not None:
|
||||
query = query.where(AuditLog.risk_score <= risk_score_max)
|
||||
|
||||
# Apply pagination
|
||||
query = query.offset(offset).limit(limit)
|
||||
query = query.order_by(AuditLog.timestamp.desc())
|
||||
|
||||
audit_logs = session.execute(query).all()
|
||||
return audit_logs
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list audit logs: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/audit-logs/{audit_id}", response_model=AgentAuditLog)
|
||||
async def get_audit_log(
|
||||
audit_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentAuditLog:
|
||||
"""Get a specific audit log entry"""
|
||||
|
||||
try:
|
||||
|
||||
audit_log = session.get(AuditLog, audit_id)
|
||||
if not audit_log:
|
||||
raise HTTPException(status_code=404, detail="Audit log not found")
|
||||
|
||||
return audit_log
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get audit log: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/trust-scores")
|
||||
async def list_trust_scores(
|
||||
entity_type: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
min_score: float | None = None,
|
||||
max_score: float | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> list[AgentTrustScore]:
|
||||
"""List trust scores with filtering"""
|
||||
|
||||
try:
|
||||
from ..services.agent_security import AgentTrustScore
|
||||
|
||||
query = select(AgentTrustScore)
|
||||
|
||||
# Apply filters
|
||||
if entity_type:
|
||||
query = query.where(AgentTrustScore.entity_type == entity_type)
|
||||
if entity_id:
|
||||
query = query.where(AgentTrustScore.entity_id == entity_id)
|
||||
if min_score is not None:
|
||||
query = query.where(AgentTrustScore.trust_score >= min_score)
|
||||
if max_score is not None:
|
||||
query = query.where(AgentTrustScore.trust_score <= max_score)
|
||||
|
||||
# Apply pagination
|
||||
query = query.offset(offset).limit(limit)
|
||||
query = query.order_by(AgentTrustScore.trust_score.desc())
|
||||
|
||||
trust_scores = session.execute(query).all()
|
||||
return trust_scores
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list trust scores: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/trust-scores/{entity_type}/{entity_id}", response_model=AgentTrustScore)
|
||||
async def get_trust_score(
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentTrustScore:
|
||||
"""Get trust score for specific entity"""
|
||||
|
||||
try:
|
||||
from ..services.agent_security import AgentTrustScore
|
||||
|
||||
trust_score = session.execute(
|
||||
select(AgentTrustScore).where(
|
||||
(AgentTrustScore.entity_type == entity_type) & (AgentTrustScore.entity_id == entity_id)
|
||||
)
|
||||
).first()
|
||||
|
||||
if not trust_score:
|
||||
raise HTTPException(status_code=404, detail="Trust score not found")
|
||||
|
||||
return trust_score
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get trust score: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/trust-scores/{entity_type}/{entity_id}/update")
|
||||
async def update_trust_score(
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
execution_success: bool,
|
||||
execution_time: float | None = None,
|
||||
security_violation: bool = False,
|
||||
policy_violation: bool = False,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> AgentTrustScore:
|
||||
"""Update trust score based on execution results"""
|
||||
|
||||
try:
|
||||
trust_manager = AgentTrustManager(session)
|
||||
trust_score = await trust_manager.update_trust_score(
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
execution_success=execution_success,
|
||||
execution_time=execution_time,
|
||||
security_violation=security_violation,
|
||||
policy_violation=policy_violation,
|
||||
)
|
||||
|
||||
# Log trust score update
|
||||
auditor = AgentAuditor(session)
|
||||
await auditor.log_event(
|
||||
AuditEventType.EXECUTION_COMPLETED if execution_success else AuditEventType.EXECUTION_FAILED,
|
||||
user_id=current_user,
|
||||
security_level=SecurityLevel.PUBLIC,
|
||||
event_data={
|
||||
"entity_type": entity_type,
|
||||
"entity_id": entity_id,
|
||||
"execution_success": execution_success,
|
||||
"execution_time": execution_time,
|
||||
"security_violation": security_violation,
|
||||
"policy_violation": policy_violation,
|
||||
},
|
||||
new_state={"trust_score": trust_score.trust_score},
|
||||
)
|
||||
|
||||
logger.info(f"Trust score updated: {entity_type}/{entity_id} -> {trust_score.trust_score}")
|
||||
return trust_score
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update trust score: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/sandbox/{execution_id}/create")
|
||||
async def create_sandbox(
|
||||
execution_id: str,
|
||||
security_level: SecurityLevel = SecurityLevel.PUBLIC,
|
||||
workflow_requirements: dict | None = None,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Create sandbox environment for agent execution"""
|
||||
|
||||
try:
|
||||
sandbox_manager = AgentSandboxManager(session)
|
||||
sandbox = await sandbox_manager.create_sandbox_environment(
|
||||
execution_id=execution_id, security_level=security_level, workflow_requirements=workflow_requirements
|
||||
)
|
||||
|
||||
# Log sandbox creation
|
||||
auditor = AgentAuditor(session)
|
||||
await auditor.log_event(
|
||||
AuditEventType.EXECUTION_STARTED,
|
||||
execution_id=execution_id,
|
||||
user_id=current_user,
|
||||
security_level=security_level,
|
||||
event_data={
|
||||
"sandbox_id": sandbox.id,
|
||||
"sandbox_type": sandbox.sandbox_type,
|
||||
"security_level": sandbox.security_level,
|
||||
},
|
||||
)
|
||||
|
||||
logger.info(f"Sandbox created for execution {execution_id}")
|
||||
return sandbox
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create sandbox: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/sandbox/{execution_id}/monitor")
|
||||
async def monitor_sandbox(
|
||||
execution_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Monitor sandbox execution for security violations"""
|
||||
|
||||
try:
|
||||
sandbox_manager = AgentSandboxManager(session)
|
||||
monitoring_data = await sandbox_manager.monitor_sandbox(execution_id)
|
||||
|
||||
return monitoring_data
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to monitor sandbox: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/sandbox/{execution_id}/cleanup")
|
||||
async def cleanup_sandbox(
|
||||
execution_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Clean up sandbox environment after execution"""
|
||||
|
||||
try:
|
||||
sandbox_manager = AgentSandboxManager(session)
|
||||
success = await sandbox_manager.cleanup_sandbox(execution_id)
|
||||
|
||||
# Log sandbox cleanup
|
||||
auditor = AgentAuditor(session)
|
||||
await auditor.log_event(
|
||||
AuditEventType.EXECUTION_COMPLETED if success else AuditEventType.EXECUTION_FAILED,
|
||||
execution_id=execution_id,
|
||||
user_id=current_user,
|
||||
security_level=SecurityLevel.PUBLIC,
|
||||
event_data={"sandbox_cleanup_success": success},
|
||||
)
|
||||
|
||||
return {"success": success, "message": "Sandbox cleanup completed"}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to cleanup sandbox: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/executions/{execution_id}/security-monitor")
|
||||
async def monitor_execution_security(
|
||||
execution_id: str,
|
||||
workflow_id: str,
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]),
|
||||
current_user: str = Depends(require_admin_key()),
|
||||
) -> dict[str, Any]:
|
||||
"""Monitor execution for security violations"""
|
||||
|
||||
try:
|
||||
security_manager = AgentSecurityManager(session)
|
||||
monitoring_result = await security_manager.monitor_execution_security(execution_id, workflow_id)
|
||||
|
||||
return monitoring_result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to monitor execution security: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/security-dashboard")
|
||||
async def get_security_dashboard(
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]), current_user: str = Depends(require_admin_key())
|
||||
) -> dict[str, Any]:
|
||||
"""Get comprehensive security dashboard data"""
|
||||
|
||||
try:
|
||||
from ..services.agent_security import AgentAuditLog, AgentSandboxConfig
|
||||
|
||||
# Get recent audit logs
|
||||
recent_audits = session.execute(select(AgentAuditLog).order_by(AgentAuditLog.timestamp.desc()).limit(50)).all()
|
||||
|
||||
# Get high-risk events
|
||||
high_risk_events = session.execute(
|
||||
select(AuditLog).where(AuditLog.requires_investigation).order_by(AuditLog.timestamp.desc()).limit(10)
|
||||
).all()
|
||||
|
||||
# Get trust score statistics
|
||||
trust_scores = session.execute(select(ActivityTrustScore)).all()
|
||||
avg_trust_score = sum(ts.trust_score for ts in trust_scores) / len(trust_scores) if trust_scores else 0
|
||||
|
||||
# Get active sandboxes
|
||||
active_sandboxes = session.execute(select(AgentSandboxConfig).where(AgentSandboxConfig.is_active)).all()
|
||||
|
||||
# Get security statistics
|
||||
total_audits = session.execute(select(AuditLog)).count()
|
||||
high_risk_count = session.execute(select(AuditLog).where(AuditLog.requires_investigation)).count()
|
||||
|
||||
security_violations = session.execute(
|
||||
select(AuditLog).where(AuditLog.event_type == AuditEventType.SECURITY_VIOLATION)
|
||||
).count()
|
||||
|
||||
return {
|
||||
"recent_audits": recent_audits,
|
||||
"high_risk_events": high_risk_events,
|
||||
"trust_score_stats": {
|
||||
"average_score": avg_trust_score,
|
||||
"total_entities": len(trust_scores),
|
||||
"high_trust_entities": len([ts for ts in trust_scores if ts.trust_score >= 80]),
|
||||
"low_trust_entities": len([ts for ts in trust_scores if ts.trust_score < 20]),
|
||||
},
|
||||
"active_sandboxes": len(active_sandboxes),
|
||||
"security_stats": {
|
||||
"total_audits": total_audits,
|
||||
"high_risk_count": high_risk_count,
|
||||
"security_violations": security_violations,
|
||||
"risk_rate": (high_risk_count / total_audits * 100) if total_audits > 0 else 0,
|
||||
},
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get security dashboard: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/security-stats")
|
||||
async def get_security_statistics(
|
||||
session: Session = Depends(Annotated[Session, Depends(get_session)]), current_user: str = Depends(require_admin_key())
|
||||
) -> dict[str, Any]:
|
||||
"""Get security statistics and metrics"""
|
||||
|
||||
try:
|
||||
from ..services.agent_security import AgentTrustScore
|
||||
|
||||
# Audit statistics
|
||||
total_audits = session.execute(select(AuditLog)).count()
|
||||
event_type_counts = {}
|
||||
for event_type in AuditEventType:
|
||||
count = session.execute(select(AuditLog).where(AuditLog.event_type == event_type)).count()
|
||||
event_type_counts[event_type.value] = count
|
||||
|
||||
# Risk score distribution
|
||||
risk_score_distribution = {"low": 0, "medium": 0, "high": 0, "critical": 0} # 0-30 # 31-70 # 71-100 # 90-100
|
||||
|
||||
all_audits = session.execute(select(AuditLog)).all()
|
||||
for audit in all_audits:
|
||||
if audit.risk_score <= 30:
|
||||
risk_score_distribution["low"] += 1
|
||||
elif audit.risk_score <= 70:
|
||||
risk_score_distribution["medium"] += 1
|
||||
elif audit.risk_score <= 90:
|
||||
risk_score_distribution["high"] += 1
|
||||
else:
|
||||
risk_score_distribution["critical"] += 1
|
||||
|
||||
# Trust score statistics
|
||||
trust_scores = session.execute(select(AgentTrustScore)).all()
|
||||
trust_score_distribution = {
|
||||
"very_low": 0, # 0-20
|
||||
"low": 0, # 21-40
|
||||
"medium": 0, # 41-60
|
||||
"high": 0, # 61-80
|
||||
"very_high": 0, # 81-100
|
||||
}
|
||||
|
||||
for trust_score in trust_scores:
|
||||
if trust_score.trust_score <= 20:
|
||||
trust_score_distribution["very_low"] += 1
|
||||
elif trust_score.trust_score <= 40:
|
||||
trust_score_distribution["low"] += 1
|
||||
elif trust_score.trust_score <= 60:
|
||||
trust_score_distribution["medium"] += 1
|
||||
elif trust_score.trust_score <= 80:
|
||||
trust_score_distribution["high"] += 1
|
||||
else:
|
||||
trust_score_distribution["very_high"] += 1
|
||||
|
||||
return {
|
||||
"audit_statistics": {
|
||||
"total_audits": total_audits,
|
||||
"event_type_counts": event_type_counts,
|
||||
"risk_score_distribution": risk_score_distribution,
|
||||
},
|
||||
"trust_statistics": {
|
||||
"total_entities": len(trust_scores),
|
||||
"average_trust_score": sum(ts.trust_score for ts in trust_scores) / len(trust_scores) if trust_scores else 0,
|
||||
"trust_score_distribution": trust_score_distribution,
|
||||
},
|
||||
"security_health": {
|
||||
"high_risk_rate": (
|
||||
(risk_score_distribution["high"] + risk_score_distribution["critical"]) / total_audits * 100
|
||||
if total_audits > 0
|
||||
else 0
|
||||
),
|
||||
"average_risk_score": sum(audit.risk_score for audit in all_audits) / len(all_audits) if all_audits else 0,
|
||||
"security_violation_rate": (
|
||||
(event_type_counts.get("security_violation", 0) / total_audits * 100) if total_audits > 0 else 0
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get security statistics: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
526
apps/agent-management/src/app/routers/services.py
Executable file
526
apps/agent-management/src/app/routers/services.py
Executable file
@@ -0,0 +1,526 @@
|
||||
from typing import Annotated
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
"""
|
||||
Services router for specific GPU workloads
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, status
|
||||
|
||||
from ..deps import require_client_key
|
||||
from ..models.services import (
|
||||
BlenderRequest,
|
||||
FFmpegRequest,
|
||||
LLMRequest,
|
||||
ServiceRequest,
|
||||
ServiceResponse,
|
||||
ServiceType,
|
||||
StableDiffusionRequest,
|
||||
WhisperRequest,
|
||||
)
|
||||
from ..schemas import JobCreate
|
||||
|
||||
# from ..models.registry import ServiceRegistry, service_registry
|
||||
from ..services import JobService
|
||||
from ..storage import get_session
|
||||
|
||||
router = APIRouter(tags=["services"])
|
||||
|
||||
|
||||
@router.post(
|
||||
"/services/{service_type}",
|
||||
response_model=ServiceResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Submit a service-specific job",
|
||||
deprecated=True,
|
||||
)
|
||||
async def submit_service_job(
|
||||
service_type: ServiceType,
|
||||
request_data: dict[str, Any],
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
user_agent: str = Header(None),
|
||||
) -> ServiceResponse:
|
||||
"""Submit a job for a specific service type
|
||||
|
||||
DEPRECATED: Use /v1/registry/services/{service_id} endpoint instead.
|
||||
This endpoint will be removed in version 2.0.
|
||||
"""
|
||||
|
||||
# Add deprecation warning header
|
||||
from fastapi import Response
|
||||
|
||||
response = Response()
|
||||
response.headers["X-Deprecated"] = "true"
|
||||
response.headers["X-Deprecation-Message"] = "Use /v1/registry/services/{service_id} instead"
|
||||
|
||||
# Check if service exists in registry
|
||||
service = service_registry.get_service(service_type.value)
|
||||
if not service:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Service {service_type} not found")
|
||||
|
||||
# Validate request against service schema
|
||||
validation_result = await validate_service_request(service_type.value, request_data)
|
||||
if not validation_result["valid"]:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid request: {', '.join(validation_result['errors'])}"
|
||||
)
|
||||
|
||||
# Create service request wrapper
|
||||
service_request = ServiceRequest(service_type=service_type, request_data=request_data)
|
||||
|
||||
# Validate and parse service-specific request
|
||||
try:
|
||||
typed_request = service_request.get_service_request()
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid request for {service_type}: {str(e)}")
|
||||
|
||||
# Get constraints from service request
|
||||
constraints = typed_request.get_constraints()
|
||||
|
||||
# Create job with service-specific payload
|
||||
job_payload = {
|
||||
"service_type": service_type.value,
|
||||
"service_request": request_data,
|
||||
}
|
||||
|
||||
job_create = JobCreate(payload=job_payload, constraints=constraints, ttl_seconds=900) # Default 15 minutes
|
||||
|
||||
# Submit job
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id, service_type=service_type, status=job.state.value, estimated_completion=job.expires_at.isoformat()
|
||||
)
|
||||
|
||||
|
||||
# Whisper endpoints
|
||||
@router.post(
|
||||
"/services/whisper/transcribe",
|
||||
response_model=ServiceResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Transcribe audio using Whisper",
|
||||
)
|
||||
async def whisper_transcribe(
|
||||
request: WhisperRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Transcribe audio file using Whisper"""
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.WHISPER.value,
|
||||
"service_request": request.dict(),
|
||||
}
|
||||
|
||||
job_create = JobCreate(payload=job_payload, constraints=request.get_constraints(), ttl_seconds=900)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.WHISPER,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/services/whisper/translate",
|
||||
response_model=ServiceResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Translate audio using Whisper",
|
||||
)
|
||||
async def whisper_translate(
|
||||
request: WhisperRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Translate audio file using Whisper"""
|
||||
# Force task to be translate
|
||||
request.task = "translate"
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.WHISPER.value,
|
||||
"service_request": request.dict(),
|
||||
}
|
||||
|
||||
job_create = JobCreate(payload=job_payload, constraints=request.get_constraints(), ttl_seconds=900)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.WHISPER,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
# Stable Diffusion endpoints
|
||||
@router.post(
|
||||
"/services/stable-diffusion/generate",
|
||||
response_model=ServiceResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Generate images using Stable Diffusion",
|
||||
)
|
||||
async def stable_diffusion_generate(
|
||||
request: StableDiffusionRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Generate images using Stable Diffusion"""
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.STABLE_DIFFUSION.value,
|
||||
"service_request": request.dict(),
|
||||
}
|
||||
|
||||
job_create = JobCreate(
|
||||
payload=job_payload, constraints=request.get_constraints(), ttl_seconds=600 # 10 minutes for image generation
|
||||
)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.STABLE_DIFFUSION,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/services/stable-diffusion/img2img",
|
||||
response_model=ServiceResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Image-to-image generation",
|
||||
)
|
||||
async def stable_diffusion_img2img(
|
||||
request: StableDiffusionRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Image-to-image generation using Stable Diffusion"""
|
||||
# Add img2img specific parameters
|
||||
request_data = request.dict()
|
||||
request_data["mode"] = "img2img"
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.STABLE_DIFFUSION.value,
|
||||
"service_request": request_data,
|
||||
}
|
||||
|
||||
job_create = JobCreate(payload=job_payload, constraints=request.get_constraints(), ttl_seconds=600)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.STABLE_DIFFUSION,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
# LLM Inference endpoints
|
||||
@router.post(
|
||||
"/services/llm/inference", response_model=ServiceResponse, status_code=status.HTTP_201_CREATED, summary="Run LLM inference"
|
||||
)
|
||||
async def llm_inference(
|
||||
request: LLMRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Run inference on a language model"""
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.LLM_INFERENCE.value,
|
||||
"service_request": request.dict(),
|
||||
}
|
||||
|
||||
job_create = JobCreate(
|
||||
payload=job_payload, constraints=request.get_constraints(), ttl_seconds=300 # 5 minutes for text generation
|
||||
)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.LLM_INFERENCE,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
@router.post("/services/llm/stream", summary="Stream LLM inference")
|
||||
async def llm_stream(
|
||||
request: LLMRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Stream LLM inference response"""
|
||||
# Force streaming mode
|
||||
request.stream = True
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.LLM_INFERENCE.value,
|
||||
"service_request": request.dict(),
|
||||
}
|
||||
|
||||
job_create = JobCreate(payload=job_payload, constraints=request.get_constraints(), ttl_seconds=300)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
# Return streaming response
|
||||
# This would implement WebSocket or Server-Sent Events
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.LLM_INFERENCE,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
# FFmpeg endpoints
|
||||
@router.post(
|
||||
"/services/ffmpeg/transcode",
|
||||
response_model=ServiceResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Transcode video using FFmpeg",
|
||||
)
|
||||
async def ffmpeg_transcode(
|
||||
request: FFmpegRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Transcode video using FFmpeg"""
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.FFMPEG.value,
|
||||
"service_request": request.dict(),
|
||||
}
|
||||
|
||||
# Adjust TTL based on video length (would need to probe video)
|
||||
job_create = JobCreate(
|
||||
payload=job_payload, constraints=request.get_constraints(), ttl_seconds=1800 # 30 minutes for video transcoding
|
||||
)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.FFMPEG,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
# Blender endpoints
|
||||
@router.post(
|
||||
"/services/blender/render",
|
||||
response_model=ServiceResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
summary="Render using Blender",
|
||||
)
|
||||
async def blender_render(
|
||||
request: BlenderRequest,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
client_id: str = Depends(require_client_key()),
|
||||
) -> ServiceResponse:
|
||||
"""Render scene using Blender"""
|
||||
|
||||
job_payload = {
|
||||
"service_type": ServiceType.BLENDER.value,
|
||||
"service_request": request.dict(),
|
||||
}
|
||||
|
||||
# Adjust TTL based on frame count
|
||||
frame_count = request.frame_end - request.frame_start + 1
|
||||
estimated_time = frame_count * 30 # 30 seconds per frame estimate
|
||||
ttl_seconds = max(600, estimated_time) # Minimum 10 minutes
|
||||
|
||||
job_create = JobCreate(payload=job_payload, constraints=request.get_constraints(), ttl_seconds=ttl_seconds)
|
||||
|
||||
service = JobService(session)
|
||||
job = service.create_job(client_id, job_create)
|
||||
|
||||
return ServiceResponse(
|
||||
job_id=job.job_id,
|
||||
service_type=ServiceType.BLENDER,
|
||||
status=job.state.value,
|
||||
estimated_completion=job.expires_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
# Utility endpoints
|
||||
@router.get("/services", summary="List available services")
|
||||
async def list_services() -> dict[str, Any]:
|
||||
"""List all available service types and their capabilities"""
|
||||
return {
|
||||
"services": [
|
||||
{
|
||||
"type": ServiceType.WHISPER.value,
|
||||
"name": "Whisper Speech Recognition",
|
||||
"description": "Transcribe and translate audio files",
|
||||
"models": [m.value for m in WhisperModel],
|
||||
"constraints": {
|
||||
"gpu": "nvidia",
|
||||
"min_vram_gb": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": ServiceType.STABLE_DIFFUSION.value,
|
||||
"name": "Stable Diffusion",
|
||||
"description": "Generate images from text prompts",
|
||||
"models": [m.value for m in SDModel],
|
||||
"constraints": {
|
||||
"gpu": "nvidia",
|
||||
"min_vram_gb": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": ServiceType.LLM_INFERENCE.value,
|
||||
"name": "LLM Inference",
|
||||
"description": "Run inference on large language models",
|
||||
"models": [m.value for m in LLMModel],
|
||||
"constraints": {
|
||||
"gpu": "nvidia",
|
||||
"min_vram_gb": 8,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": ServiceType.FFMPEG.value,
|
||||
"name": "FFmpeg Video Processing",
|
||||
"description": "Transcode and process video files",
|
||||
"codecs": [c.value for c in FFmpegCodec],
|
||||
"constraints": {
|
||||
"gpu": "any",
|
||||
"min_vram_gb": 0,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": ServiceType.BLENDER.value,
|
||||
"name": "Blender Rendering",
|
||||
"description": "Render 3D scenes using Blender",
|
||||
"engines": [e.value for e in BlenderEngine],
|
||||
"constraints": {
|
||||
"gpu": "any",
|
||||
"min_vram_gb": 4,
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@router.get("/services/{service_type}/schema", summary="Get service request schema", deprecated=True)
|
||||
async def get_service_schema(service_type: ServiceType) -> dict[str, Any]:
|
||||
"""Get the JSON schema for a specific service type
|
||||
|
||||
DEPRECATED: Use /v1/registry/services/{service_id}/schema instead.
|
||||
This endpoint will be removed in version 2.0.
|
||||
"""
|
||||
# Get service from registry
|
||||
service = service_registry.get_service(service_type.value)
|
||||
if not service:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Service {service_type} not found")
|
||||
|
||||
# Build schema from service definition
|
||||
properties = {}
|
||||
required = []
|
||||
|
||||
for param in service.input_parameters:
|
||||
prop = {"type": param.type.value, "description": param.description}
|
||||
|
||||
if param.default is not None:
|
||||
prop["default"] = param.default
|
||||
if param.min_value is not None:
|
||||
prop["minimum"] = param.min_value
|
||||
if param.max_value is not None:
|
||||
prop["maximum"] = param.max_value
|
||||
if param.options:
|
||||
prop["enum"] = param.options
|
||||
if param.validation:
|
||||
prop.update(param.validation)
|
||||
|
||||
properties[param.name] = prop
|
||||
if param.required:
|
||||
required.append(param.name)
|
||||
|
||||
schema = {"type": "object", "properties": properties, "required": required}
|
||||
|
||||
return {"service_type": service_type.value, "schema": schema}
|
||||
|
||||
|
||||
async def validate_service_request(service_id: str, request_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Validate a service request against the service schema"""
|
||||
service = service_registry.get_service(service_id)
|
||||
if not service:
|
||||
return {"valid": False, "errors": [f"Service {service_id} not found"]}
|
||||
|
||||
validation_result = {"valid": True, "errors": [], "warnings": []}
|
||||
|
||||
# Check required parameters
|
||||
provided_params = set(request_data.keys())
|
||||
required_params = {p.name for p in service.input_parameters if p.required}
|
||||
missing_params = required_params - provided_params
|
||||
|
||||
if missing_params:
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].extend([f"Missing required parameter: {param}" for param in missing_params])
|
||||
|
||||
# Validate parameter types and constraints
|
||||
for param in service.input_parameters:
|
||||
if param.name in request_data:
|
||||
value = request_data[param.name]
|
||||
|
||||
# Type validation (simplified)
|
||||
if param.type == "integer" and not isinstance(value, int):
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].append(f"Parameter {param.name} must be an integer")
|
||||
elif param.type == "float" and not isinstance(value, (int, float)):
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].append(f"Parameter {param.name} must be a number")
|
||||
elif param.type == "boolean" and not isinstance(value, bool):
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].append(f"Parameter {param.name} must be a boolean")
|
||||
elif param.type == "array" and not isinstance(value, list):
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].append(f"Parameter {param.name} must be an array")
|
||||
|
||||
# Value constraints
|
||||
if param.min_value is not None and value < param.min_value:
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].append(f"Parameter {param.name} must be >= {param.min_value}")
|
||||
|
||||
if param.max_value is not None and value > param.max_value:
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].append(f"Parameter {param.name} must be <= {param.max_value}")
|
||||
|
||||
# Enum options
|
||||
if param.options and value not in param.options:
|
||||
validation_result["valid"] = False
|
||||
validation_result["errors"].append(f"Parameter {param.name} must be one of: {', '.join(param.options)}")
|
||||
|
||||
return validation_result
|
||||
|
||||
|
||||
# Import models for type hints
|
||||
from ..models.services import (
|
||||
BlenderEngine,
|
||||
FFmpegCodec,
|
||||
LLMModel,
|
||||
SDModel,
|
||||
WhisperModel,
|
||||
)
|
||||
0
apps/agent-management/src/app/services/__init__.py
Normal file
0
apps/agent-management/src/app/services/__init__.py
Normal file
102
apps/agent-management/src/app/services/advanced_rl/agents.py
Normal file
102
apps/agent-management/src/app/services/advanced_rl/agents.py
Normal file
@@ -0,0 +1,102 @@
|
||||
"""
|
||||
Reinforcement Learning Agent Models
|
||||
PyTorch neural network models for various RL algorithms
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class PPOAgent(nn.Module):
|
||||
"""Proximal Policy Optimization Agent"""
|
||||
|
||||
def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 256):
|
||||
super().__init__()
|
||||
self.actor = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, action_dim),
|
||||
nn.Softmax(dim=-1),
|
||||
)
|
||||
self.critic = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1)
|
||||
)
|
||||
|
||||
def forward(self, state):
|
||||
action_probs = self.actor(state)
|
||||
value = self.critic(state)
|
||||
return action_probs, value
|
||||
|
||||
|
||||
class SACAgent(nn.Module):
|
||||
"""Soft Actor-Critic Agent"""
|
||||
|
||||
def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 256):
|
||||
super().__init__()
|
||||
self.actor_mean = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, action_dim),
|
||||
)
|
||||
self.actor_log_std = nn.Parameter(torch.zeros(1, action_dim))
|
||||
|
||||
self.qf1 = nn.Sequential(
|
||||
nn.Linear(state_dim + action_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, 1),
|
||||
)
|
||||
|
||||
self.qf2 = nn.Sequential(
|
||||
nn.Linear(state_dim + action_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, 1),
|
||||
)
|
||||
|
||||
def forward(self, state):
|
||||
mean = self.actor_mean(state)
|
||||
std = torch.exp(self.actor_log_std)
|
||||
return mean, std
|
||||
|
||||
|
||||
class RainbowDQNAgent(nn.Module):
|
||||
"""Rainbow DQN Agent with multiple improvements"""
|
||||
|
||||
def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 512, num_atoms: int = 51):
|
||||
super().__init__()
|
||||
self.num_atoms = num_atoms
|
||||
self.action_dim = action_dim
|
||||
|
||||
# Feature extractor
|
||||
self.feature_layer = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU()
|
||||
)
|
||||
|
||||
# Dueling network architecture
|
||||
self.value_stream = nn.Sequential(
|
||||
nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, num_atoms)
|
||||
)
|
||||
|
||||
self.advantage_stream = nn.Sequential(
|
||||
nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, action_dim * num_atoms)
|
||||
)
|
||||
|
||||
def forward(self, state):
|
||||
features = self.feature_layer(state)
|
||||
values = self.value_stream(features)
|
||||
advantages = self.advantage_stream(features)
|
||||
|
||||
# Reshape for distributional RL
|
||||
advantages = advantages.view(-1, self.action_dim, self.num_atoms)
|
||||
values = values.view(-1, 1, self.num_atoms)
|
||||
|
||||
# Dueling architecture
|
||||
q_atoms = values + advantages - advantages.mean(dim=1, keepdim=True)
|
||||
return q_atoms
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
PPO Agent implementation
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class PPOAgent(nn.Module):
|
||||
"""Proximal Policy Optimization Agent"""
|
||||
|
||||
def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 256):
|
||||
super().__init__()
|
||||
self.actor = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, action_dim),
|
||||
nn.Softmax(dim=-1),
|
||||
)
|
||||
self.critic = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1)
|
||||
)
|
||||
|
||||
def forward(self, state):
|
||||
action_probs = self.actor(state)
|
||||
value = self.critic(state)
|
||||
return action_probs, value
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
Rainbow DQN Agent implementation
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class RainbowDQNAgent(nn.Module):
|
||||
"""Rainbow DQN Agent with multiple improvements"""
|
||||
|
||||
def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 512, num_atoms: int = 51):
|
||||
super().__init__()
|
||||
self.num_atoms = num_atoms
|
||||
self.action_dim = action_dim
|
||||
|
||||
# Feature extractor
|
||||
self.feature_layer = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU()
|
||||
)
|
||||
|
||||
# Dueling network architecture
|
||||
self.value_stream = nn.Sequential(
|
||||
nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, num_atoms)
|
||||
)
|
||||
|
||||
self.advantage_stream = nn.Sequential(
|
||||
nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, action_dim * num_atoms)
|
||||
)
|
||||
|
||||
def forward(self, state):
|
||||
features = self.feature_layer(state)
|
||||
values = self.value_stream(features)
|
||||
advantages = self.advantage_stream(features)
|
||||
|
||||
# Reshape for distributional RL
|
||||
advantages = advantages.view(-1, self.action_dim, self.num_atoms)
|
||||
values = values.view(-1, 1, self.num_atoms)
|
||||
|
||||
# Dueling architecture
|
||||
q_atoms = values + advantages - advantages.mean(dim=1, keepdim=True)
|
||||
return q_atoms
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
SAC Agent implementation
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class SACAgent(nn.Module):
|
||||
"""Soft Actor-Critic Agent"""
|
||||
|
||||
def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 256):
|
||||
super().__init__()
|
||||
self.actor_mean = nn.Sequential(
|
||||
nn.Linear(state_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, action_dim),
|
||||
)
|
||||
self.actor_log_std = nn.Parameter(torch.zeros(1, action_dim))
|
||||
|
||||
self.qf1 = nn.Sequential(
|
||||
nn.Linear(state_dim + action_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, 1),
|
||||
)
|
||||
|
||||
self.qf2 = nn.Sequential(
|
||||
nn.Linear(state_dim + action_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(hidden_dim, 1),
|
||||
)
|
||||
|
||||
def forward(self, state):
|
||||
mean = self.actor_mean(state)
|
||||
std = torch.exp(self.actor_log_std)
|
||||
return mean, std
|
||||
988
apps/agent-management/src/app/services/agent_communication.py
Executable file
988
apps/agent-management/src/app/services/agent_communication.py
Executable file
@@ -0,0 +1,988 @@
|
||||
"""
|
||||
Agent Communication Service for Advanced Agent Features
|
||||
Implements secure agent-to-agent messaging with reputation-based access control
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from .cross_chain_reputation import CrossChainReputationService
|
||||
|
||||
|
||||
class MessageType(StrEnum):
|
||||
"""Types of agent messages"""
|
||||
|
||||
TEXT = "text"
|
||||
DATA = "data"
|
||||
TASK_REQUEST = "task_request"
|
||||
TASK_RESPONSE = "task_response"
|
||||
COLLABORATION = "collaboration"
|
||||
NOTIFICATION = "notification"
|
||||
SYSTEM = "system"
|
||||
URGENT = "urgent"
|
||||
BULK = "bulk"
|
||||
|
||||
|
||||
class ChannelType(StrEnum):
|
||||
"""Types of communication channels"""
|
||||
|
||||
DIRECT = "direct"
|
||||
GROUP = "group"
|
||||
BROADCAST = "broadcast"
|
||||
PRIVATE = "private"
|
||||
|
||||
|
||||
class MessageStatus(StrEnum):
|
||||
"""Message delivery status"""
|
||||
|
||||
PENDING = "pending"
|
||||
DELIVERED = "delivered"
|
||||
READ = "read"
|
||||
FAILED = "failed"
|
||||
EXPIRED = "expired"
|
||||
|
||||
|
||||
class EncryptionType(StrEnum):
|
||||
"""Encryption types for messages"""
|
||||
|
||||
AES256 = "aes256"
|
||||
RSA = "rsa"
|
||||
HYBRID = "hybrid"
|
||||
NONE = "none"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Message:
|
||||
"""Agent message data"""
|
||||
|
||||
id: str
|
||||
sender: str
|
||||
recipient: str
|
||||
message_type: MessageType
|
||||
content: bytes
|
||||
encryption_key: bytes
|
||||
encryption_type: EncryptionType
|
||||
size: int
|
||||
timestamp: datetime
|
||||
delivery_timestamp: datetime | None = None
|
||||
read_timestamp: datetime | None = None
|
||||
status: MessageStatus = MessageStatus.PENDING
|
||||
paid: bool = False
|
||||
price: float = 0.0
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
expires_at: datetime | None = None
|
||||
reply_to: str | None = None
|
||||
thread_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class CommunicationChannel:
|
||||
"""Communication channel between agents"""
|
||||
|
||||
id: str
|
||||
agent1: str
|
||||
agent2: str
|
||||
channel_type: ChannelType
|
||||
is_active: bool
|
||||
created_timestamp: datetime
|
||||
last_activity: datetime
|
||||
message_count: int
|
||||
participants: list[str] = field(default_factory=list)
|
||||
encryption_enabled: bool = True
|
||||
auto_delete: bool = False
|
||||
retention_period: int = 2592000 # 30 days
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageTemplate:
|
||||
"""Message template for common communications"""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
description: str
|
||||
message_type: MessageType
|
||||
content_template: str
|
||||
variables: list[str]
|
||||
base_price: float
|
||||
is_active: bool
|
||||
creator: str
|
||||
usage_count: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class CommunicationStats:
|
||||
"""Communication statistics for agent"""
|
||||
|
||||
total_messages: int
|
||||
total_earnings: float
|
||||
messages_sent: int
|
||||
messages_received: int
|
||||
active_channels: int
|
||||
last_activity: datetime
|
||||
average_response_time: float
|
||||
delivery_rate: float
|
||||
|
||||
|
||||
class AgentCommunicationService:
|
||||
"""Service for managing agent-to-agent communication"""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
self.config = config
|
||||
self.messages: dict[str, Message] = {}
|
||||
self.channels: dict[str, CommunicationChannel] = {}
|
||||
self.message_templates: dict[str, MessageTemplate] = {}
|
||||
self.agent_messages: dict[str, list[str]] = {}
|
||||
self.agent_channels: dict[str, list[str]] = {}
|
||||
self.communication_stats: dict[str, CommunicationStats] = {}
|
||||
|
||||
# Services
|
||||
self.reputation_service: CrossChainReputationService | None = None
|
||||
|
||||
# Configuration
|
||||
self.min_reputation_score = 1000
|
||||
self.base_message_price = 0.001 # AITBC
|
||||
self.max_message_size = 100000 # 100KB
|
||||
self.message_timeout = 86400 # 24 hours
|
||||
self.channel_timeout = 2592000 # 30 days
|
||||
self.encryption_enabled = True
|
||||
|
||||
# Access control
|
||||
self.authorized_agents: dict[str, bool] = {}
|
||||
self.contact_lists: dict[str, dict[str, bool]] = {}
|
||||
self.blocked_lists: dict[str, dict[str, bool]] = {}
|
||||
|
||||
# Message routing
|
||||
self.message_queue: list[Message] = []
|
||||
self.delivery_attempts: dict[str, int] = {}
|
||||
|
||||
# Templates
|
||||
self._initialize_default_templates()
|
||||
|
||||
def set_reputation_service(self, reputation_service: CrossChainReputationService):
|
||||
"""Set reputation service for access control"""
|
||||
self.reputation_service = reputation_service
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize the agent communication service"""
|
||||
logger.info("Initializing Agent Communication Service")
|
||||
|
||||
# Load existing data
|
||||
await self._load_communication_data()
|
||||
|
||||
# Start background tasks
|
||||
asyncio.create_task(self._process_message_queue())
|
||||
asyncio.create_task(self._cleanup_expired_messages())
|
||||
asyncio.create_task(self._cleanup_inactive_channels())
|
||||
|
||||
logger.info("Agent Communication Service initialized")
|
||||
|
||||
async def authorize_agent(self, agent_id: str) -> bool:
|
||||
"""Authorize an agent to use the communication system"""
|
||||
|
||||
try:
|
||||
self.authorized_agents[agent_id] = True
|
||||
|
||||
# Initialize communication stats
|
||||
if agent_id not in self.communication_stats:
|
||||
self.communication_stats[agent_id] = CommunicationStats(
|
||||
total_messages=0,
|
||||
total_earnings=0.0,
|
||||
messages_sent=0,
|
||||
messages_received=0,
|
||||
active_channels=0,
|
||||
last_activity=datetime.now(timezone.utc),
|
||||
average_response_time=0.0,
|
||||
delivery_rate=0.0,
|
||||
)
|
||||
|
||||
logger.info(f"Authorized agent: {agent_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to authorize agent {agent_id}: {e}")
|
||||
return False
|
||||
|
||||
async def revoke_agent(self, agent_id: str) -> bool:
|
||||
"""Revoke agent authorization"""
|
||||
|
||||
try:
|
||||
self.authorized_agents[agent_id] = False
|
||||
|
||||
# Clean up agent data
|
||||
if agent_id in self.agent_messages:
|
||||
del self.agent_messages[agent_id]
|
||||
if agent_id in self.agent_channels:
|
||||
del self.agent_channels[agent_id]
|
||||
if agent_id in self.communication_stats:
|
||||
del self.communication_stats[agent_id]
|
||||
|
||||
logger.info(f"Revoked authorization for agent: {agent_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to revoke agent {agent_id}: {e}")
|
||||
return False
|
||||
|
||||
async def add_contact(self, agent_id: str, contact_id: str) -> bool:
|
||||
"""Add contact to agent's contact list"""
|
||||
|
||||
try:
|
||||
if agent_id not in self.contact_lists:
|
||||
self.contact_lists[agent_id] = {}
|
||||
|
||||
self.contact_lists[agent_id][contact_id] = True
|
||||
|
||||
# Remove from blocked list if present
|
||||
if agent_id in self.blocked_lists and contact_id in self.blocked_lists[agent_id]:
|
||||
del self.blocked_lists[agent_id][contact_id]
|
||||
|
||||
logger.info(f"Added contact {contact_id} for agent {agent_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to add contact: {e}")
|
||||
return False
|
||||
|
||||
async def remove_contact(self, agent_id: str, contact_id: str) -> bool:
|
||||
"""Remove contact from agent's contact list"""
|
||||
|
||||
try:
|
||||
if agent_id in self.contact_lists and contact_id in self.contact_lists[agent_id]:
|
||||
del self.contact_lists[agent_id][contact_id]
|
||||
|
||||
logger.info(f"Removed contact {contact_id} for agent {agent_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to remove contact: {e}")
|
||||
return False
|
||||
|
||||
async def block_agent(self, agent_id: str, blocked_id: str) -> bool:
|
||||
"""Block an agent"""
|
||||
|
||||
try:
|
||||
if agent_id not in self.blocked_lists:
|
||||
self.blocked_lists[agent_id] = {}
|
||||
|
||||
self.blocked_lists[agent_id][blocked_id] = True
|
||||
|
||||
# Remove from contact list if present
|
||||
if agent_id in self.contact_lists and blocked_id in self.contact_lists[agent_id]:
|
||||
del self.contact_lists[agent_id][blocked_id]
|
||||
|
||||
logger.info(f"Blocked agent {blocked_id} for agent {agent_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to block agent: {e}")
|
||||
return False
|
||||
|
||||
async def unblock_agent(self, agent_id: str, blocked_id: str) -> bool:
|
||||
"""Unblock an agent"""
|
||||
|
||||
try:
|
||||
if agent_id in self.blocked_lists and blocked_id in self.blocked_lists[agent_id]:
|
||||
del self.blocked_lists[agent_id][blocked_id]
|
||||
|
||||
logger.info(f"Unblocked agent {blocked_id} for agent {agent_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to unblock agent: {e}")
|
||||
return False
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
sender: str,
|
||||
recipient: str,
|
||||
message_type: MessageType,
|
||||
content: str,
|
||||
encryption_type: EncryptionType = EncryptionType.AES256,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
reply_to: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
) -> str:
|
||||
"""Send a message to another agent"""
|
||||
|
||||
try:
|
||||
# Validate authorization
|
||||
if not await self._can_send_message(sender, recipient):
|
||||
raise PermissionError("Not authorized to send message")
|
||||
|
||||
# Validate content
|
||||
content_bytes = content.encode("utf-8")
|
||||
if len(content_bytes) > self.max_message_size:
|
||||
raise ValueError(f"Message too large: {len(content_bytes)} > {self.max_message_size}")
|
||||
|
||||
# Generate message ID
|
||||
message_id = await self._generate_message_id()
|
||||
|
||||
# Encrypt content
|
||||
if encryption_type != EncryptionType.NONE:
|
||||
encrypted_content, encryption_key = await self._encrypt_content(content_bytes, encryption_type)
|
||||
else:
|
||||
encrypted_content = content_bytes
|
||||
encryption_key = b""
|
||||
|
||||
# Calculate price
|
||||
price = await self._calculate_message_price(len(content_bytes), message_type)
|
||||
|
||||
# Create message
|
||||
message = Message(
|
||||
id=message_id,
|
||||
sender=sender,
|
||||
recipient=recipient,
|
||||
message_type=message_type,
|
||||
content=encrypted_content,
|
||||
encryption_key=encryption_key,
|
||||
encryption_type=encryption_type,
|
||||
size=len(content_bytes),
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
status=MessageStatus.PENDING,
|
||||
price=price,
|
||||
metadata=metadata or {},
|
||||
expires_at=datetime.now(timezone.utc) + timedelta(seconds=self.message_timeout),
|
||||
reply_to=reply_to,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
|
||||
# Store message
|
||||
self.messages[message_id] = message
|
||||
|
||||
# Update message lists
|
||||
if sender not in self.agent_messages:
|
||||
self.agent_messages[sender] = []
|
||||
if recipient not in self.agent_messages:
|
||||
self.agent_messages[recipient] = []
|
||||
|
||||
self.agent_messages[sender].append(message_id)
|
||||
self.agent_messages[recipient].append(message_id)
|
||||
|
||||
# Update stats
|
||||
await self._update_message_stats(sender, recipient, "sent")
|
||||
|
||||
# Create or update channel
|
||||
await self._get_or_create_channel(sender, recipient, ChannelType.DIRECT)
|
||||
|
||||
# Add to queue for delivery
|
||||
self.message_queue.append(message)
|
||||
|
||||
logger.info(f"Message sent from {sender} to {recipient}: {message_id}")
|
||||
return message_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message: {e}")
|
||||
raise
|
||||
|
||||
async def deliver_message(self, message_id: str) -> bool:
|
||||
"""Mark message as delivered"""
|
||||
|
||||
try:
|
||||
if message_id not in self.messages:
|
||||
raise ValueError(f"Message {message_id} not found")
|
||||
|
||||
message = self.messages[message_id]
|
||||
if message.status != MessageStatus.PENDING:
|
||||
raise ValueError(f"Message {message_id} not pending")
|
||||
|
||||
message.status = MessageStatus.DELIVERED
|
||||
message.delivery_timestamp = datetime.now(timezone.utc)
|
||||
|
||||
# Update stats
|
||||
await self._update_message_stats(message.sender, message.recipient, "delivered")
|
||||
|
||||
logger.info(f"Message delivered: {message_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to deliver message {message_id}: {e}")
|
||||
return False
|
||||
|
||||
async def read_message(self, message_id: str, reader: str) -> str | None:
|
||||
"""Mark message as read and return decrypted content"""
|
||||
|
||||
try:
|
||||
if message_id not in self.messages:
|
||||
raise ValueError(f"Message {message_id} not found")
|
||||
|
||||
message = self.messages[message_id]
|
||||
if message.recipient != reader:
|
||||
raise PermissionError("Not message recipient")
|
||||
|
||||
if message.status != MessageStatus.DELIVERED:
|
||||
raise ValueError("Message not delivered")
|
||||
|
||||
if message.read:
|
||||
raise ValueError("Message already read")
|
||||
|
||||
# Mark as read
|
||||
message.status = MessageStatus.READ
|
||||
message.read_timestamp = datetime.now(timezone.utc)
|
||||
|
||||
# Update stats
|
||||
await self._update_message_stats(message.sender, message.recipient, "read")
|
||||
|
||||
# Decrypt content
|
||||
if message.encryption_type != EncryptionType.NONE:
|
||||
decrypted_content = await self._decrypt_content(
|
||||
message.content, message.encryption_key, message.encryption_type
|
||||
)
|
||||
return decrypted_content.decode("utf-8")
|
||||
else:
|
||||
return message.content.decode("utf-8")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to read message {message_id}: {e}")
|
||||
return None
|
||||
|
||||
async def pay_for_message(self, message_id: str, payer: str, amount: float) -> bool:
|
||||
"""Pay for a message"""
|
||||
|
||||
try:
|
||||
if message_id not in self.messages:
|
||||
raise ValueError(f"Message {message_id} not found")
|
||||
|
||||
message = self.messages[message_id]
|
||||
|
||||
if amount < message.price:
|
||||
raise ValueError(f"Insufficient payment: {amount} < {message.price}")
|
||||
|
||||
# Process payment (simplified)
|
||||
# In production, implement actual payment processing
|
||||
|
||||
message.paid = True
|
||||
|
||||
# Update sender's earnings
|
||||
if message.sender in self.communication_stats:
|
||||
self.communication_stats[message.sender].total_earnings += message.price
|
||||
|
||||
logger.info(f"Payment processed for message {message_id}: {amount}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process payment for message {message_id}: {e}")
|
||||
return False
|
||||
|
||||
async def create_channel(
|
||||
self, agent1: str, agent2: str, channel_type: ChannelType = ChannelType.DIRECT, encryption_enabled: bool = True
|
||||
) -> str:
|
||||
"""Create a communication channel"""
|
||||
|
||||
try:
|
||||
# Validate agents
|
||||
if not self.authorized_agents.get(agent1, False) or not self.authorized_agents.get(agent2, False):
|
||||
raise PermissionError("Agents not authorized")
|
||||
|
||||
if agent1 == agent2:
|
||||
raise ValueError("Cannot create channel with self")
|
||||
|
||||
# Generate channel ID
|
||||
channel_id = await self._generate_channel_id()
|
||||
|
||||
# Create channel
|
||||
channel = CommunicationChannel(
|
||||
id=channel_id,
|
||||
agent1=agent1,
|
||||
agent2=agent2,
|
||||
channel_type=channel_type,
|
||||
is_active=True,
|
||||
created_timestamp=datetime.now(timezone.utc),
|
||||
last_activity=datetime.now(timezone.utc),
|
||||
message_count=0,
|
||||
participants=[agent1, agent2],
|
||||
encryption_enabled=encryption_enabled,
|
||||
)
|
||||
|
||||
# Store channel
|
||||
self.channels[channel_id] = channel
|
||||
|
||||
# Update agent channel lists
|
||||
if agent1 not in self.agent_channels:
|
||||
self.agent_channels[agent1] = []
|
||||
if agent2 not in self.agent_channels:
|
||||
self.agent_channels[agent2] = []
|
||||
|
||||
self.agent_channels[agent1].append(channel_id)
|
||||
self.agent_channels[agent2].append(channel_id)
|
||||
|
||||
# Update stats
|
||||
self.communication_stats[agent1].active_channels += 1
|
||||
self.communication_stats[agent2].active_channels += 1
|
||||
|
||||
logger.info(f"Channel created: {channel_id} between {agent1} and {agent2}")
|
||||
return channel_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create channel: {e}")
|
||||
raise
|
||||
|
||||
async def create_message_template(
|
||||
self,
|
||||
creator: str,
|
||||
name: str,
|
||||
description: str,
|
||||
message_type: MessageType,
|
||||
content_template: str,
|
||||
variables: list[str],
|
||||
base_price: float = 0.001,
|
||||
) -> str:
|
||||
"""Create a message template"""
|
||||
|
||||
try:
|
||||
# Generate template ID
|
||||
template_id = await self._generate_template_id()
|
||||
|
||||
template = MessageTemplate(
|
||||
id=template_id,
|
||||
name=name,
|
||||
description=description,
|
||||
message_type=message_type,
|
||||
content_template=content_template,
|
||||
variables=variables,
|
||||
base_price=base_price,
|
||||
is_active=True,
|
||||
creator=creator,
|
||||
)
|
||||
|
||||
self.message_templates[template_id] = template
|
||||
|
||||
logger.info(f"Template created: {template_id}")
|
||||
return template_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create template: {e}")
|
||||
raise
|
||||
|
||||
async def use_template(self, template_id: str, sender: str, recipient: str, variables: dict[str, str]) -> str:
|
||||
"""Use a message template to send a message"""
|
||||
|
||||
try:
|
||||
if template_id not in self.message_templates:
|
||||
raise ValueError(f"Template {template_id} not found")
|
||||
|
||||
template = self.message_templates[template_id]
|
||||
|
||||
if not template.is_active:
|
||||
raise ValueError(f"Template {template_id} not active")
|
||||
|
||||
# Substitute variables
|
||||
content = template.content_template
|
||||
for var, value in variables.items():
|
||||
if var in template.variables:
|
||||
content = content.replace(f"{{{var}}}", value)
|
||||
|
||||
# Send message
|
||||
message_id = await self.send_message(
|
||||
sender=sender,
|
||||
recipient=recipient,
|
||||
message_type=template.message_type,
|
||||
content=content,
|
||||
metadata={"template_id": template_id},
|
||||
)
|
||||
|
||||
# Update template usage
|
||||
template.usage_count += 1
|
||||
|
||||
logger.info(f"Template used: {template_id} -> {message_id}")
|
||||
return message_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to use template {template_id}: {e}")
|
||||
raise
|
||||
|
||||
async def get_agent_messages(
|
||||
self, agent_id: str, limit: int = 50, offset: int = 0, status: MessageStatus | None = None
|
||||
) -> list[Message]:
|
||||
"""Get messages for an agent"""
|
||||
|
||||
try:
|
||||
if agent_id not in self.agent_messages:
|
||||
return []
|
||||
|
||||
message_ids = self.agent_messages[agent_id]
|
||||
|
||||
# Apply filters
|
||||
filtered_messages = []
|
||||
for message_id in message_ids:
|
||||
if message_id in self.messages:
|
||||
message = self.messages[message_id]
|
||||
if status is None or message.status == status:
|
||||
filtered_messages.append(message)
|
||||
|
||||
# Sort by timestamp (newest first)
|
||||
filtered_messages.sort(key=lambda x: x.timestamp, reverse=True)
|
||||
|
||||
# Apply pagination
|
||||
return filtered_messages[offset : offset + limit]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get messages for {agent_id}: {e}")
|
||||
return []
|
||||
|
||||
async def get_unread_messages(self, agent_id: str) -> list[Message]:
|
||||
"""Get unread messages for an agent"""
|
||||
|
||||
try:
|
||||
if agent_id not in self.agent_messages:
|
||||
return []
|
||||
|
||||
unread_messages = []
|
||||
for message_id in self.agent_messages[agent_id]:
|
||||
if message_id in self.messages:
|
||||
message = self.messages[message_id]
|
||||
if message.recipient == agent_id and message.status == MessageStatus.DELIVERED:
|
||||
unread_messages.append(message)
|
||||
|
||||
return unread_messages
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get unread messages for {agent_id}: {e}")
|
||||
return []
|
||||
|
||||
async def get_agent_channels(self, agent_id: str) -> list[CommunicationChannel]:
|
||||
"""Get channels for an agent"""
|
||||
|
||||
try:
|
||||
if agent_id not in self.agent_channels:
|
||||
return []
|
||||
|
||||
channels = []
|
||||
for channel_id in self.agent_channels[agent_id]:
|
||||
if channel_id in self.channels:
|
||||
channels.append(self.channels[channel_id])
|
||||
|
||||
return channels
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get channels for {agent_id}: {e}")
|
||||
return []
|
||||
|
||||
async def get_communication_stats(self, agent_id: str) -> CommunicationStats:
|
||||
"""Get communication statistics for an agent"""
|
||||
|
||||
try:
|
||||
if agent_id not in self.communication_stats:
|
||||
raise ValueError(f"Agent {agent_id} not found")
|
||||
|
||||
return self.communication_stats[agent_id]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get stats for {agent_id}: {e}")
|
||||
raise
|
||||
|
||||
async def can_communicate(self, sender: str, recipient: str) -> bool:
|
||||
"""Check if agents can communicate"""
|
||||
|
||||
# Check authorization
|
||||
if not self.authorized_agents.get(sender, False) or not self.authorized_agents.get(recipient, False):
|
||||
return False
|
||||
|
||||
# Check blocked lists
|
||||
if (sender in self.blocked_lists and recipient in self.blocked_lists[sender]) or (
|
||||
recipient in self.blocked_lists and sender in self.blocked_lists[recipient]
|
||||
):
|
||||
return False
|
||||
|
||||
# Check contact lists
|
||||
if sender in self.contact_lists and recipient in self.contact_lists[sender]:
|
||||
return True
|
||||
|
||||
# Check reputation
|
||||
if self.reputation_service:
|
||||
sender_reputation = await self.reputation_service.get_reputation_score(sender)
|
||||
return sender_reputation >= self.min_reputation_score
|
||||
|
||||
return False
|
||||
|
||||
async def _can_send_message(self, sender: str, recipient: str) -> bool:
|
||||
"""Check if sender can send message to recipient"""
|
||||
return await self.can_communicate(sender, recipient)
|
||||
|
||||
async def _generate_message_id(self) -> str:
|
||||
"""Generate unique message ID"""
|
||||
import uuid
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
async def _generate_channel_id(self) -> str:
|
||||
"""Generate unique channel ID"""
|
||||
import uuid
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
async def _generate_template_id(self) -> str:
|
||||
"""Generate unique template ID"""
|
||||
import uuid
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
async def _encrypt_content(self, content: bytes, encryption_type: EncryptionType) -> tuple[bytes, bytes]:
|
||||
"""Encrypt message content"""
|
||||
|
||||
if encryption_type == EncryptionType.AES256:
|
||||
# Simplified AES encryption
|
||||
key = hashlib.sha256(content).digest()[:32] # Generate key from content
|
||||
import os
|
||||
|
||||
iv = os.urandom(16)
|
||||
|
||||
# In production, use proper AES encryption
|
||||
encrypted = content + iv # Simplified
|
||||
return encrypted, key
|
||||
|
||||
elif encryption_type == EncryptionType.RSA:
|
||||
# Simplified RSA encryption
|
||||
key = hashlib.sha256(content).digest()[:256]
|
||||
return content + key, key
|
||||
|
||||
else:
|
||||
return content, b""
|
||||
|
||||
async def _decrypt_content(self, encrypted_content: bytes, key: bytes, encryption_type: EncryptionType) -> bytes:
|
||||
"""Decrypt message content"""
|
||||
|
||||
if encryption_type == EncryptionType.AES256:
|
||||
# Simplified AES decryption
|
||||
if len(encrypted_content) < 16:
|
||||
return encrypted_content
|
||||
return encrypted_content[:-16] # Remove IV
|
||||
|
||||
elif encryption_type == EncryptionType.RSA:
|
||||
# Simplified RSA decryption
|
||||
if len(encrypted_content) < 256:
|
||||
return encrypted_content
|
||||
return encrypted_content[:-256] # Remove key
|
||||
|
||||
else:
|
||||
return encrypted_content
|
||||
|
||||
async def _calculate_message_price(self, size: int, message_type: MessageType) -> float:
|
||||
"""Calculate message price based on size and type"""
|
||||
|
||||
base_price = self.base_message_price
|
||||
|
||||
# Size multiplier
|
||||
size_multiplier = max(1, size / 1000) # 1 AITBC per 1000 bytes
|
||||
|
||||
# Type multiplier
|
||||
type_multipliers = {
|
||||
MessageType.TEXT: 1.0,
|
||||
MessageType.DATA: 1.5,
|
||||
MessageType.TASK_REQUEST: 2.0,
|
||||
MessageType.TASK_RESPONSE: 2.0,
|
||||
MessageType.COLLABORATION: 3.0,
|
||||
MessageType.NOTIFICATION: 0.5,
|
||||
MessageType.SYSTEM: 0.1,
|
||||
MessageType.URGENT: 5.0,
|
||||
MessageType.BULK: 10.0,
|
||||
}
|
||||
|
||||
type_multiplier = type_multipliers.get(message_type, 1.0)
|
||||
|
||||
return base_price * size_multiplier * type_multiplier
|
||||
|
||||
async def _get_or_create_channel(self, agent1: str, agent2: str, channel_type: ChannelType) -> str:
|
||||
"""Get or create communication channel"""
|
||||
|
||||
# Check if channel already exists
|
||||
if agent1 in self.agent_channels:
|
||||
for channel_id in self.agent_channels[agent1]:
|
||||
if channel_id in self.channels:
|
||||
channel = self.channels[channel_id]
|
||||
if channel.is_active and (
|
||||
(channel.agent1 == agent1 and channel.agent2 == agent2)
|
||||
or (channel.agent1 == agent2 and channel.agent2 == agent1)
|
||||
):
|
||||
return channel_id
|
||||
|
||||
# Create new channel
|
||||
return await self.create_channel(agent1, agent2, channel_type)
|
||||
|
||||
async def _update_message_stats(self, sender: str, recipient: str, action: str):
|
||||
"""Update message statistics"""
|
||||
|
||||
if action == "sent":
|
||||
if sender in self.communication_stats:
|
||||
self.communication_stats[sender].total_messages += 1
|
||||
self.communication_stats[sender].messages_sent += 1
|
||||
self.communication_stats[sender].last_activity = datetime.now(timezone.utc)
|
||||
|
||||
elif action == "delivered":
|
||||
if recipient in self.communication_stats:
|
||||
self.communication_stats[recipient].total_messages += 1
|
||||
self.communication_stats[recipient].messages_received += 1
|
||||
self.communication_stats[recipient].last_activity = datetime.now(timezone.utc)
|
||||
|
||||
elif action == "read":
|
||||
if recipient in self.communication_stats:
|
||||
self.communication_stats[recipient].last_activity = datetime.now(timezone.utc)
|
||||
|
||||
async def _process_message_queue(self):
|
||||
"""Process message queue for delivery"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
if self.message_queue:
|
||||
message = self.message_queue.pop(0)
|
||||
|
||||
# Simulate delivery
|
||||
await asyncio.sleep(0.1)
|
||||
await self.deliver_message(message.id)
|
||||
|
||||
await asyncio.sleep(1)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing message queue: {e}")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
async def _cleanup_expired_messages(self):
|
||||
"""Clean up expired messages"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
current_time = datetime.now(timezone.utc)
|
||||
expired_messages = []
|
||||
|
||||
for message_id, message in self.messages.items():
|
||||
if message.expires_at and current_time > message.expires_at:
|
||||
expired_messages.append(message_id)
|
||||
|
||||
for message_id in expired_messages:
|
||||
del self.messages[message_id]
|
||||
# Remove from agent message lists
|
||||
for _agent_id, message_ids in self.agent_messages.items():
|
||||
if message_id in message_ids:
|
||||
message_ids.remove(message_id)
|
||||
|
||||
if expired_messages:
|
||||
logger.info(f"Cleaned up {len(expired_messages)} expired messages")
|
||||
|
||||
await asyncio.sleep(3600) # Check every hour
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up messages: {e}")
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
async def _cleanup_inactive_channels(self):
|
||||
"""Clean up inactive channels"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
current_time = datetime.now(timezone.utc)
|
||||
inactive_channels = []
|
||||
|
||||
for channel_id, channel in self.channels.items():
|
||||
if channel.is_active and current_time > channel.last_activity + timedelta(seconds=self.channel_timeout):
|
||||
inactive_channels.append(channel_id)
|
||||
|
||||
for channel_id in inactive_channels:
|
||||
channel = self.channels[channel_id]
|
||||
channel.is_active = False
|
||||
|
||||
# Update stats
|
||||
if channel.agent1 in self.communication_stats:
|
||||
self.communication_stats[channel.agent1].active_channels = max(
|
||||
0, self.communication_stats[channel.agent1].active_channels - 1
|
||||
)
|
||||
if channel.agent2 in self.communication_stats:
|
||||
self.communication_stats[channel.agent2].active_channels = max(
|
||||
0, self.communication_stats[channel.agent2].active_channels - 1
|
||||
)
|
||||
|
||||
if inactive_channels:
|
||||
logger.info(f"Cleaned up {len(inactive_channels)} inactive channels")
|
||||
|
||||
await asyncio.sleep(3600) # Check every hour
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up channels: {e}")
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
def _initialize_default_templates(self):
|
||||
"""Initialize default message templates"""
|
||||
|
||||
templates = [
|
||||
MessageTemplate(
|
||||
id="task_request_default",
|
||||
name="Task Request",
|
||||
description="Default template for task requests",
|
||||
message_type=MessageType.TASK_REQUEST,
|
||||
content_template="Hello! I have a task for you: {task_description}. Budget: {budget} AITBC. Deadline: {deadline}.",
|
||||
variables=["task_description", "budget", "deadline"],
|
||||
base_price=0.002,
|
||||
is_active=True,
|
||||
creator="system",
|
||||
),
|
||||
MessageTemplate(
|
||||
id="collaboration_invite",
|
||||
name="Collaboration Invite",
|
||||
description="Template for inviting agents to collaborate",
|
||||
message_type=MessageType.COLLABORATION,
|
||||
content_template="I'd like to collaborate on {project_name}. Your role would be {role_description}. Interested?",
|
||||
variables=["project_name", "role_description"],
|
||||
base_price=0.003,
|
||||
is_active=True,
|
||||
creator="system",
|
||||
),
|
||||
MessageTemplate(
|
||||
id="notification_update",
|
||||
name="Notification Update",
|
||||
description="Template for sending notifications",
|
||||
message_type=MessageType.NOTIFICATION,
|
||||
content_template="Notification: {notification_type}. {message}. Action required: {action_required}.",
|
||||
variables=["notification_type", "message", "action_required"],
|
||||
base_price=0.001,
|
||||
is_active=True,
|
||||
creator="system",
|
||||
),
|
||||
]
|
||||
|
||||
for template in templates:
|
||||
self.message_templates[template.id] = template
|
||||
|
||||
async def _load_communication_data(self):
|
||||
"""Load existing communication data"""
|
||||
# In production, load from database
|
||||
pass
|
||||
|
||||
async def export_communication_data(self, format: str = "json") -> str:
|
||||
"""Export communication data"""
|
||||
|
||||
data = {
|
||||
"messages": {k: asdict(v) for k, v in self.messages.items()},
|
||||
"channels": {k: asdict(v) for k, v in self.channels.items()},
|
||||
"templates": {k: asdict(v) for k, v in self.message_templates.items()},
|
||||
"export_timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
|
||||
if format.lower() == "json":
|
||||
return json.dumps(data, indent=2, default=str)
|
||||
else:
|
||||
raise ValueError(f"Unsupported format: {format}")
|
||||
|
||||
async def import_communication_data(self, data: str, format: str = "json"):
|
||||
"""Import communication data"""
|
||||
|
||||
if format.lower() == "json":
|
||||
parsed_data = json.loads(data)
|
||||
|
||||
# Import messages
|
||||
for message_id, message_data in parsed_data.get("messages", {}).items():
|
||||
message_data["timestamp"] = datetime.fromisoformat(message_data["timestamp"])
|
||||
self.messages[message_id] = Message(**message_data)
|
||||
|
||||
# Import channels
|
||||
for channel_id, channel_data in parsed_data.get("channels", {}).items():
|
||||
channel_data["created_timestamp"] = datetime.fromisoformat(channel_data["created_timestamp"])
|
||||
channel_data["last_activity"] = datetime.fromisoformat(channel_data["last_activity"])
|
||||
self.channels[channel_id] = CommunicationChannel(**channel_data)
|
||||
|
||||
logger.info("Communication data imported successfully")
|
||||
else:
|
||||
raise ValueError(f"Unsupported format: {format}")
|
||||
1159
apps/agent-management/src/app/services/agent_integration.py
Executable file
1159
apps/agent-management/src/app/services/agent_integration.py
Executable file
File diff suppressed because it is too large
Load Diff
692
apps/agent-management/src/app/services/agent_orchestrator.py
Executable file
692
apps/agent-management/src/app/services/agent_orchestrator.py
Executable file
@@ -0,0 +1,692 @@
|
||||
"""
|
||||
Agent Orchestrator Service for hermes Autonomous Economics
|
||||
Implements multi-agent coordination and sub-task management
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from .bid_strategy_engine import BidResult
|
||||
from .task_decomposition import GPU_Tier, SubTask, SubTaskStatus, TaskDecomposition
|
||||
|
||||
|
||||
class OrchestratorStatus(StrEnum):
|
||||
"""Orchestrator status"""
|
||||
|
||||
IDLE = "idle"
|
||||
PLANNING = "planning"
|
||||
EXECUTING = "executing"
|
||||
MONITORING = "monitoring"
|
||||
FAILED = "failed"
|
||||
COMPLETED = "completed"
|
||||
|
||||
|
||||
class AgentStatus(StrEnum):
|
||||
"""Agent status"""
|
||||
|
||||
AVAILABLE = "available"
|
||||
BUSY = "busy"
|
||||
OFFLINE = "offline"
|
||||
MAINTENANCE = "maintenance"
|
||||
|
||||
|
||||
class ResourceType(StrEnum):
|
||||
"""Resource types"""
|
||||
|
||||
GPU = "gpu"
|
||||
CPU = "cpu"
|
||||
MEMORY = "memory"
|
||||
STORAGE = "storage"
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentCapability:
|
||||
"""Agent capability definition"""
|
||||
|
||||
agent_id: str
|
||||
supported_task_types: list[str]
|
||||
gpu_tier: GPU_Tier
|
||||
max_concurrent_tasks: int
|
||||
current_load: int
|
||||
performance_score: float # 0-1
|
||||
cost_per_hour: float
|
||||
reliability_score: float # 0-1
|
||||
last_updated: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResourceAllocation:
|
||||
"""Resource allocation for an agent"""
|
||||
|
||||
agent_id: str
|
||||
sub_task_id: str
|
||||
resource_type: ResourceType
|
||||
allocated_amount: int
|
||||
allocated_at: datetime
|
||||
expected_duration: float
|
||||
actual_duration: float | None = None
|
||||
cost: float | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentAssignment:
|
||||
"""Assignment of sub-task to agent"""
|
||||
|
||||
sub_task_id: str
|
||||
agent_id: str
|
||||
assigned_at: datetime
|
||||
started_at: datetime | None = None
|
||||
completed_at: datetime | None = None
|
||||
status: SubTaskStatus = SubTaskStatus.PENDING
|
||||
bid_result: BidResult | None = None
|
||||
resource_allocations: list[ResourceAllocation] = field(default_factory=list)
|
||||
error_message: str | None = None
|
||||
retry_count: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrchestrationPlan:
|
||||
"""Complete orchestration plan for a task"""
|
||||
|
||||
task_id: str
|
||||
decomposition: TaskDecomposition
|
||||
agent_assignments: list[AgentAssignment]
|
||||
execution_timeline: dict[str, datetime]
|
||||
resource_requirements: dict[ResourceType, int]
|
||||
estimated_cost: float
|
||||
confidence_score: float
|
||||
created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
class AgentOrchestrator:
|
||||
"""Multi-agent orchestration service"""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
self.config = config
|
||||
self.status = OrchestratorStatus.IDLE
|
||||
|
||||
# Agent registry
|
||||
self.agent_capabilities: dict[str, AgentCapability] = {}
|
||||
self.agent_status: dict[str, AgentStatus] = {}
|
||||
|
||||
# Orchestration tracking
|
||||
self.active_plans: dict[str, OrchestrationPlan] = {}
|
||||
self.completed_plans: list[OrchestrationPlan] = []
|
||||
self.failed_plans: list[OrchestrationPlan] = []
|
||||
|
||||
# Resource tracking
|
||||
self.resource_allocations: dict[str, list[ResourceAllocation]] = {}
|
||||
self.resource_utilization: dict[ResourceType, float] = {}
|
||||
|
||||
# Performance metrics
|
||||
self.orchestration_metrics = {
|
||||
"total_tasks": 0,
|
||||
"successful_tasks": 0,
|
||||
"failed_tasks": 0,
|
||||
"average_execution_time": 0.0,
|
||||
"average_cost": 0.0,
|
||||
"agent_utilization": 0.0,
|
||||
}
|
||||
|
||||
# Configuration
|
||||
self.max_concurrent_plans = config.get("max_concurrent_plans", 10)
|
||||
self.assignment_timeout = config.get("assignment_timeout", 300) # 5 minutes
|
||||
self.monitoring_interval = config.get("monitoring_interval", 30) # 30 seconds
|
||||
self.retry_limit = config.get("retry_limit", 3)
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize the orchestrator"""
|
||||
logger.info("Initializing Agent Orchestrator")
|
||||
|
||||
# Load agent capabilities
|
||||
await self._load_agent_capabilities()
|
||||
|
||||
# Start monitoring
|
||||
asyncio.create_task(self._monitor_executions())
|
||||
asyncio.create_task(self._update_agent_status())
|
||||
|
||||
logger.info("Agent Orchestrator initialized")
|
||||
|
||||
async def orchestrate_task(
|
||||
self,
|
||||
task_id: str,
|
||||
decomposition: TaskDecomposition,
|
||||
budget_limit: float | None = None,
|
||||
deadline: datetime | None = None,
|
||||
) -> OrchestrationPlan:
|
||||
"""Orchestrate execution of a decomposed task"""
|
||||
|
||||
try:
|
||||
logger.info(f"Orchestrating task {task_id} with {len(decomposition.sub_tasks)} sub-tasks")
|
||||
|
||||
# Check capacity
|
||||
if len(self.active_plans) >= self.max_concurrent_plans:
|
||||
raise Exception("Orchestrator at maximum capacity")
|
||||
|
||||
self.status = OrchestratorStatus.PLANNING
|
||||
|
||||
# Create orchestration plan
|
||||
plan = await self._create_orchestration_plan(task_id, decomposition, budget_limit, deadline)
|
||||
|
||||
# Execute assignments
|
||||
await self._execute_assignments(plan)
|
||||
|
||||
# Start monitoring
|
||||
self.active_plans[task_id] = plan
|
||||
self.status = OrchestratorStatus.MONITORING
|
||||
|
||||
# Update metrics
|
||||
self.orchestration_metrics["total_tasks"] += 1
|
||||
|
||||
logger.info(f"Task {task_id} orchestration plan created and started")
|
||||
return plan
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to orchestrate task {task_id}: {e}")
|
||||
self.status = OrchestratorStatus.FAILED
|
||||
raise
|
||||
|
||||
async def get_task_status(self, task_id: str) -> dict[str, Any]:
|
||||
"""Get status of orchestrated task"""
|
||||
|
||||
if task_id not in self.active_plans:
|
||||
return {"status": "not_found"}
|
||||
|
||||
plan = self.active_plans[task_id]
|
||||
|
||||
# Count sub-task statuses
|
||||
status_counts = {}
|
||||
for status in SubTaskStatus:
|
||||
status_counts[status.value] = 0
|
||||
|
||||
completed_count = 0
|
||||
failed_count = 0
|
||||
|
||||
for assignment in plan.agent_assignments:
|
||||
status_counts[assignment.status.value] += 1
|
||||
|
||||
if assignment.status == SubTaskStatus.COMPLETED:
|
||||
completed_count += 1
|
||||
elif assignment.status == SubTaskStatus.FAILED:
|
||||
failed_count += 1
|
||||
|
||||
# Determine overall status
|
||||
total_sub_tasks = len(plan.agent_assignments)
|
||||
if completed_count == total_sub_tasks:
|
||||
overall_status = "completed"
|
||||
elif failed_count > 0:
|
||||
overall_status = "failed"
|
||||
elif completed_count > 0:
|
||||
overall_status = "in_progress"
|
||||
else:
|
||||
overall_status = "pending"
|
||||
|
||||
return {
|
||||
"status": overall_status,
|
||||
"progress": completed_count / total_sub_tasks if total_sub_tasks > 0 else 0,
|
||||
"completed_sub_tasks": completed_count,
|
||||
"failed_sub_tasks": failed_count,
|
||||
"total_sub_tasks": total_sub_tasks,
|
||||
"estimated_cost": plan.estimated_cost,
|
||||
"actual_cost": await self._calculate_actual_cost(plan),
|
||||
"started_at": plan.created_at.isoformat(),
|
||||
"assignments": [
|
||||
{
|
||||
"sub_task_id": a.sub_task_id,
|
||||
"agent_id": a.agent_id,
|
||||
"status": a.status.value,
|
||||
"assigned_at": a.assigned_at.isoformat(),
|
||||
"started_at": a.started_at.isoformat() if a.started_at else None,
|
||||
"completed_at": a.completed_at.isoformat() if a.completed_at else None,
|
||||
}
|
||||
for a in plan.agent_assignments
|
||||
],
|
||||
}
|
||||
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
"""Cancel task orchestration"""
|
||||
|
||||
if task_id not in self.active_plans:
|
||||
return False
|
||||
|
||||
plan = self.active_plans[task_id]
|
||||
|
||||
# Cancel all active assignments
|
||||
for assignment in plan.agent_assignments:
|
||||
if assignment.status in [SubTaskStatus.PENDING, SubTaskStatus.IN_PROGRESS]:
|
||||
assignment.status = SubTaskStatus.CANCELLED
|
||||
await self._release_agent_resources(assignment.agent_id, assignment.sub_task_id)
|
||||
|
||||
# Move to failed plans
|
||||
self.failed_plans.append(plan)
|
||||
del self.active_plans[task_id]
|
||||
|
||||
logger.info(f"Task {task_id} cancelled")
|
||||
return True
|
||||
|
||||
async def retry_failed_sub_tasks(self, task_id: str) -> list[str]:
|
||||
"""Retry failed sub-tasks"""
|
||||
|
||||
if task_id not in self.active_plans:
|
||||
return []
|
||||
|
||||
plan = self.active_plans[task_id]
|
||||
retried_tasks = []
|
||||
|
||||
for assignment in plan.agent_assignments:
|
||||
if assignment.status == SubTaskStatus.FAILED and assignment.retry_count < self.retry_limit:
|
||||
# Reset assignment
|
||||
assignment.status = SubTaskStatus.PENDING
|
||||
assignment.started_at = None
|
||||
assignment.completed_at = None
|
||||
assignment.error_message = None
|
||||
assignment.retry_count += 1
|
||||
|
||||
# Release resources
|
||||
await self._release_agent_resources(assignment.agent_id, assignment.sub_task_id)
|
||||
|
||||
# Re-assign
|
||||
await self._assign_sub_task(assignment.sub_task_id, plan)
|
||||
|
||||
retried_tasks.append(assignment.sub_task_id)
|
||||
logger.info(f"Retrying sub-task {assignment.sub_task_id} (attempt {assignment.retry_count + 1})")
|
||||
|
||||
return retried_tasks
|
||||
|
||||
async def register_agent(self, capability: AgentCapability):
|
||||
"""Register a new agent"""
|
||||
|
||||
self.agent_capabilities[capability.agent_id] = capability
|
||||
self.agent_status[capability.agent_id] = AgentStatus.AVAILABLE
|
||||
|
||||
logger.info(f"Registered agent {capability.agent_id}")
|
||||
|
||||
async def update_agent_status(self, agent_id: str, status: AgentStatus):
|
||||
"""Update agent status"""
|
||||
|
||||
if agent_id in self.agent_status:
|
||||
self.agent_status[agent_id] = status
|
||||
logger.info(f"Updated agent {agent_id} status to {status}")
|
||||
|
||||
async def get_available_agents(self, task_type: str, gpu_tier: GPU_Tier) -> list[AgentCapability]:
|
||||
"""Get available agents for task"""
|
||||
|
||||
available_agents = []
|
||||
|
||||
for agent_id, capability in self.agent_capabilities.items():
|
||||
if (
|
||||
self.agent_status.get(agent_id) == AgentStatus.AVAILABLE
|
||||
and task_type in capability.supported_task_types
|
||||
and capability.gpu_tier == gpu_tier
|
||||
and capability.current_load < capability.max_concurrent_tasks
|
||||
):
|
||||
available_agents.append(capability)
|
||||
|
||||
# Sort by performance score
|
||||
available_agents.sort(key=lambda x: x.performance_score, reverse=True)
|
||||
|
||||
return available_agents
|
||||
|
||||
async def get_orchestration_metrics(self) -> dict[str, Any]:
|
||||
"""Get orchestration performance metrics"""
|
||||
|
||||
return {
|
||||
"orchestrator_status": self.status.value,
|
||||
"active_plans": len(self.active_plans),
|
||||
"completed_plans": len(self.completed_plans),
|
||||
"failed_plans": len(self.failed_plans),
|
||||
"registered_agents": len(self.agent_capabilities),
|
||||
"available_agents": len([s for s in self.agent_status.values() if s == AgentStatus.AVAILABLE]),
|
||||
"metrics": self.orchestration_metrics,
|
||||
"resource_utilization": self.resource_utilization,
|
||||
}
|
||||
|
||||
async def _create_orchestration_plan(
|
||||
self, task_id: str, decomposition: TaskDecomposition, budget_limit: float | None, deadline: datetime | None
|
||||
) -> OrchestrationPlan:
|
||||
"""Create detailed orchestration plan"""
|
||||
|
||||
assignments = []
|
||||
execution_timeline = {}
|
||||
resource_requirements = dict.fromkeys(ResourceType, 0)
|
||||
total_cost = 0.0
|
||||
|
||||
# Process each execution stage
|
||||
for stage_idx, stage_sub_tasks in enumerate(decomposition.execution_plan):
|
||||
stage_start = datetime.now(timezone.utc) + timedelta(hours=stage_idx * 2) # Estimate 2 hours per stage
|
||||
|
||||
for sub_task_id in stage_sub_tasks:
|
||||
# Find sub-task
|
||||
sub_task = next(st for st in decomposition.sub_tasks if st.sub_task_id == sub_task_id)
|
||||
|
||||
# Create assignment (will be filled during execution)
|
||||
assignment = AgentAssignment(
|
||||
sub_task_id=sub_task_id, agent_id="", assigned_at=datetime.now(timezone.utc) # Will be assigned during execution
|
||||
)
|
||||
assignments.append(assignment)
|
||||
|
||||
# Calculate resource requirements
|
||||
resource_requirements[ResourceType.GPU] += 1
|
||||
resource_requirements[ResourceType.MEMORY] += sub_task.requirements.memory_requirement
|
||||
|
||||
# Set timeline
|
||||
execution_timeline[sub_task_id] = stage_start
|
||||
|
||||
# Calculate confidence score
|
||||
confidence_score = await self._calculate_plan_confidence(decomposition, budget_limit, deadline)
|
||||
|
||||
return OrchestrationPlan(
|
||||
task_id=task_id,
|
||||
decomposition=decomposition,
|
||||
agent_assignments=assignments,
|
||||
execution_timeline=execution_timeline,
|
||||
resource_requirements=resource_requirements,
|
||||
estimated_cost=total_cost,
|
||||
confidence_score=confidence_score,
|
||||
)
|
||||
|
||||
async def _execute_assignments(self, plan: OrchestrationPlan):
|
||||
"""Execute agent assignments"""
|
||||
|
||||
for assignment in plan.agent_assignments:
|
||||
await self._assign_sub_task(assignment.sub_task_id, plan)
|
||||
|
||||
async def _assign_sub_task(self, sub_task_id: str, plan: OrchestrationPlan):
|
||||
"""Assign sub-task to suitable agent"""
|
||||
|
||||
# Find sub-task
|
||||
sub_task = next(st for st in plan.decomposition.sub_tasks if st.sub_task_id == sub_task_id)
|
||||
|
||||
# Get available agents
|
||||
available_agents = await self.get_available_agents(
|
||||
sub_task.requirements.task_type.value, sub_task.requirements.gpu_tier
|
||||
)
|
||||
|
||||
if not available_agents:
|
||||
raise Exception(f"No available agents for sub-task {sub_task_id}")
|
||||
|
||||
# Select best agent
|
||||
best_agent = await self._select_best_agent(available_agents, sub_task)
|
||||
|
||||
# Update assignment
|
||||
assignment = next(a for a in plan.agent_assignments if a.sub_task_id == sub_task_id)
|
||||
assignment.agent_id = best_agent.agent_id
|
||||
assignment.status = SubTaskStatus.ASSIGNED
|
||||
|
||||
# Update agent load
|
||||
self.agent_capabilities[best_agent.agent_id].current_load += 1
|
||||
self.agent_status[best_agent.agent_id] = AgentStatus.BUSY
|
||||
|
||||
# Allocate resources
|
||||
await self._allocate_resources(best_agent.agent_id, sub_task_id, sub_task.requirements)
|
||||
|
||||
logger.info(f"Assigned sub-task {sub_task_id} to agent {best_agent.agent_id}")
|
||||
|
||||
async def _select_best_agent(self, available_agents: list[AgentCapability], sub_task: SubTask) -> AgentCapability:
|
||||
"""Select best agent for sub-task"""
|
||||
|
||||
# Score agents based on multiple factors
|
||||
scored_agents = []
|
||||
|
||||
for agent in available_agents:
|
||||
score = 0.0
|
||||
|
||||
# Performance score (40% weight)
|
||||
score += agent.performance_score * 0.4
|
||||
|
||||
# Cost efficiency (30% weight)
|
||||
cost_efficiency = min(1.0, 0.05 / agent.cost_per_hour) # Normalize around 0.05 AITBC/hour
|
||||
score += cost_efficiency * 0.3
|
||||
|
||||
# Reliability (20% weight)
|
||||
score += agent.reliability_score * 0.2
|
||||
|
||||
# Current load (10% weight)
|
||||
load_factor = 1.0 - (agent.current_load / agent.max_concurrent_tasks)
|
||||
score += load_factor * 0.1
|
||||
|
||||
scored_agents.append((agent, score))
|
||||
|
||||
# Select highest scoring agent
|
||||
scored_agents.sort(key=lambda x: x[1], reverse=True)
|
||||
return scored_agents[0][0]
|
||||
|
||||
async def _allocate_resources(self, agent_id: str, sub_task_id: str, requirements):
|
||||
"""Allocate resources for sub-task"""
|
||||
|
||||
allocations = []
|
||||
|
||||
# GPU allocation
|
||||
gpu_allocation = ResourceAllocation(
|
||||
agent_id=agent_id,
|
||||
sub_task_id=sub_task_id,
|
||||
resource_type=ResourceType.GPU,
|
||||
allocated_amount=1,
|
||||
allocated_at=datetime.now(timezone.utc),
|
||||
expected_duration=requirements.estimated_duration,
|
||||
)
|
||||
allocations.append(gpu_allocation)
|
||||
|
||||
# Memory allocation
|
||||
memory_allocation = ResourceAllocation(
|
||||
agent_id=agent_id,
|
||||
sub_task_id=sub_task_id,
|
||||
resource_type=ResourceType.MEMORY,
|
||||
allocated_amount=requirements.memory_requirement,
|
||||
allocated_at=datetime.now(timezone.utc),
|
||||
expected_duration=requirements.estimated_duration,
|
||||
)
|
||||
allocations.append(memory_allocation)
|
||||
|
||||
# Store allocations
|
||||
if agent_id not in self.resource_allocations:
|
||||
self.resource_allocations[agent_id] = []
|
||||
self.resource_allocations[agent_id].extend(allocations)
|
||||
|
||||
async def _release_agent_resources(self, agent_id: str, sub_task_id: str):
|
||||
"""Release resources from agent"""
|
||||
|
||||
if agent_id in self.resource_allocations:
|
||||
# Remove allocations for this sub-task
|
||||
self.resource_allocations[agent_id] = [
|
||||
alloc for alloc in self.resource_allocations[agent_id] if alloc.sub_task_id != sub_task_id
|
||||
]
|
||||
|
||||
# Update agent load
|
||||
if agent_id in self.agent_capabilities:
|
||||
self.agent_capabilities[agent_id].current_load = max(0, self.agent_capabilities[agent_id].current_load - 1)
|
||||
|
||||
# Update status if no load
|
||||
if self.agent_capabilities[agent_id].current_load == 0:
|
||||
self.agent_status[agent_id] = AgentStatus.AVAILABLE
|
||||
|
||||
async def _monitor_executions(self):
|
||||
"""Monitor active executions"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Check all active plans
|
||||
completed_tasks = []
|
||||
failed_tasks = []
|
||||
|
||||
for task_id, plan in list(self.active_plans.items()):
|
||||
# Check if all sub-tasks are completed
|
||||
all_completed = all(a.status == SubTaskStatus.COMPLETED for a in plan.agent_assignments)
|
||||
any_failed = any(a.status == SubTaskStatus.FAILED for a in plan.agent_assignments)
|
||||
|
||||
if all_completed:
|
||||
completed_tasks.append(task_id)
|
||||
elif any_failed:
|
||||
# Check if all failed tasks have exceeded retry limit
|
||||
all_failed_exhausted = all(
|
||||
a.status == SubTaskStatus.FAILED and a.retry_count >= self.retry_limit
|
||||
for a in plan.agent_assignments
|
||||
if a.status == SubTaskStatus.FAILED
|
||||
)
|
||||
if all_failed_exhausted:
|
||||
failed_tasks.append(task_id)
|
||||
|
||||
# Move completed/failed tasks
|
||||
for task_id in completed_tasks:
|
||||
plan = self.active_plans[task_id]
|
||||
self.completed_plans.append(plan)
|
||||
del self.active_plans[task_id]
|
||||
self.orchestration_metrics["successful_tasks"] += 1
|
||||
logger.info(f"Task {task_id} completed successfully")
|
||||
|
||||
for task_id in failed_tasks:
|
||||
plan = self.active_plans[task_id]
|
||||
self.failed_plans.append(plan)
|
||||
del self.active_plans[task_id]
|
||||
self.orchestration_metrics["failed_tasks"] += 1
|
||||
logger.info(f"Task {task_id} failed")
|
||||
|
||||
# Update resource utilization
|
||||
await self._update_resource_utilization()
|
||||
|
||||
await asyncio.sleep(self.monitoring_interval)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in execution monitoring: {e}")
|
||||
await asyncio.sleep(60)
|
||||
|
||||
async def _update_agent_status(self):
|
||||
"""Update agent status periodically"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Check agent health and update status
|
||||
for agent_id in self.agent_capabilities.keys():
|
||||
# In a real implementation, this would ping agents or check health endpoints
|
||||
# For now, assume agents are healthy if they have recent updates
|
||||
|
||||
capability = self.agent_capabilities[agent_id]
|
||||
time_since_update = datetime.now(timezone.utc) - capability.last_updated
|
||||
|
||||
if time_since_update > timedelta(minutes=5):
|
||||
if self.agent_status[agent_id] != AgentStatus.OFFLINE:
|
||||
self.agent_status[agent_id] = AgentStatus.OFFLINE
|
||||
logger.warning(f"Agent {agent_id} marked as offline")
|
||||
elif self.agent_status[agent_id] == AgentStatus.OFFLINE:
|
||||
self.agent_status[agent_id] = AgentStatus.AVAILABLE
|
||||
logger.info(f"Agent {agent_id} back online")
|
||||
|
||||
await asyncio.sleep(60) # Check every minute
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating agent status: {e}")
|
||||
await asyncio.sleep(60)
|
||||
|
||||
async def _update_resource_utilization(self):
|
||||
"""Update resource utilization metrics"""
|
||||
|
||||
total_resources = dict.fromkeys(ResourceType, 0)
|
||||
used_resources = dict.fromkeys(ResourceType, 0)
|
||||
|
||||
# Calculate total resources
|
||||
for capability in self.agent_capabilities.values():
|
||||
total_resources[ResourceType.GPU] += capability.max_concurrent_tasks
|
||||
# Add other resource types as needed
|
||||
|
||||
# Calculate used resources
|
||||
for allocations in self.resource_allocations.values():
|
||||
for allocation in allocations:
|
||||
used_resources[allocation.resource_type] += allocation.allocated_amount
|
||||
|
||||
# Calculate utilization
|
||||
for resource_type in ResourceType:
|
||||
total = total_resources[resource_type]
|
||||
used = used_resources[resource_type]
|
||||
self.resource_utilization[resource_type] = used / total if total > 0 else 0.0
|
||||
|
||||
async def _calculate_plan_confidence(
|
||||
self, decomposition: TaskDecomposition, budget_limit: float | None, deadline: datetime | None
|
||||
) -> float:
|
||||
"""Calculate confidence in orchestration plan"""
|
||||
|
||||
confidence = decomposition.confidence_score
|
||||
|
||||
# Adjust for budget constraints
|
||||
if budget_limit and decomposition.estimated_total_cost > budget_limit:
|
||||
confidence *= 0.7
|
||||
|
||||
# Adjust for deadline
|
||||
if deadline:
|
||||
time_to_deadline = (deadline - datetime.now(timezone.utc)).total_seconds() / 3600
|
||||
if time_to_deadline < decomposition.estimated_total_duration:
|
||||
confidence *= 0.6
|
||||
|
||||
# Adjust for agent availability
|
||||
available_agents = len([s for s in self.agent_status.values() if s == AgentStatus.AVAILABLE])
|
||||
total_agents = len(self.agent_capabilities)
|
||||
|
||||
if total_agents > 0:
|
||||
availability_ratio = available_agents / total_agents
|
||||
confidence *= 0.5 + availability_ratio * 0.5
|
||||
|
||||
return max(0.1, min(0.95, confidence))
|
||||
|
||||
async def _calculate_actual_cost(self, plan: OrchestrationPlan) -> float:
|
||||
"""Calculate actual cost of orchestration"""
|
||||
|
||||
actual_cost = 0.0
|
||||
|
||||
for assignment in plan.agent_assignments:
|
||||
if assignment.agent_id in self.agent_capabilities:
|
||||
agent = self.agent_capabilities[assignment.agent_id]
|
||||
|
||||
# Calculate cost based on actual duration
|
||||
duration = assignment.actual_duration or 1.0 # Default to 1 hour
|
||||
cost = agent.cost_per_hour * duration
|
||||
actual_cost += cost
|
||||
|
||||
return actual_cost
|
||||
|
||||
async def _load_agent_capabilities(self):
|
||||
"""Load agent capabilities from storage"""
|
||||
|
||||
# In a real implementation, this would load from database or configuration
|
||||
# For now, create some mock agents
|
||||
|
||||
mock_agents = [
|
||||
AgentCapability(
|
||||
agent_id="agent_001",
|
||||
supported_task_types=["text_processing", "data_analysis"],
|
||||
gpu_tier=GPU_Tier.MID_RANGE_GPU,
|
||||
max_concurrent_tasks=3,
|
||||
current_load=0,
|
||||
performance_score=0.85,
|
||||
cost_per_hour=0.05,
|
||||
reliability_score=0.92,
|
||||
),
|
||||
AgentCapability(
|
||||
agent_id="agent_002",
|
||||
supported_task_types=["image_processing", "model_inference"],
|
||||
gpu_tier=GPU_Tier.HIGH_END_GPU,
|
||||
max_concurrent_tasks=2,
|
||||
current_load=0,
|
||||
performance_score=0.92,
|
||||
cost_per_hour=0.09,
|
||||
reliability_score=0.88,
|
||||
),
|
||||
AgentCapability(
|
||||
agent_id="agent_003",
|
||||
supported_task_types=["compute_intensive", "model_training"],
|
||||
gpu_tier=GPU_Tier.PREMIUM_GPU,
|
||||
max_concurrent_tasks=1,
|
||||
current_load=0,
|
||||
performance_score=0.96,
|
||||
cost_per_hour=0.15,
|
||||
reliability_score=0.95,
|
||||
),
|
||||
]
|
||||
|
||||
for agent in mock_agents:
|
||||
await self.register_agent(agent)
|
||||
988
apps/agent-management/src/app/services/agent_performance_service.py
Executable file
988
apps/agent-management/src/app/services/agent_performance_service.py
Executable file
@@ -0,0 +1,988 @@
|
||||
"""
|
||||
Advanced Agent Performance Service
|
||||
Implements meta-learning, resource optimization, and performance enhancement for hermes agents
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from app.domain.agent_performance import (
|
||||
AgentPerformanceProfile,
|
||||
LearningStrategy,
|
||||
MetaLearningModel,
|
||||
OptimizationTarget,
|
||||
PerformanceMetric,
|
||||
PerformanceOptimization,
|
||||
ResourceAllocation,
|
||||
ResourceType,
|
||||
)
|
||||
|
||||
|
||||
class MetaLearningEngine:
|
||||
"""Advanced meta-learning system for rapid skill acquisition"""
|
||||
|
||||
def __init__(self):
|
||||
self.meta_algorithms = {
|
||||
"model_agnostic_meta_learning": self.maml_algorithm,
|
||||
"reptile": self.reptile_algorithm,
|
||||
"meta_sgd": self.meta_sgd_algorithm,
|
||||
"prototypical_networks": self.prototypical_algorithm,
|
||||
}
|
||||
|
||||
self.adaptation_strategies = {
|
||||
"fast_adaptation": self.fast_adaptation,
|
||||
"gradual_adaptation": self.gradual_adaptation,
|
||||
"transfer_adaptation": self.transfer_adaptation,
|
||||
"multi_task_adaptation": self.multi_task_adaptation,
|
||||
}
|
||||
|
||||
self.performance_metrics = [
|
||||
PerformanceMetric.ACCURACY,
|
||||
PerformanceMetric.ADAPTATION_SPEED,
|
||||
PerformanceMetric.GENERALIZATION,
|
||||
PerformanceMetric.RESOURCE_EFFICIENCY,
|
||||
]
|
||||
|
||||
async def create_meta_learning_model(
|
||||
self,
|
||||
session: Session,
|
||||
model_name: str,
|
||||
base_algorithms: list[str],
|
||||
meta_strategy: LearningStrategy,
|
||||
adaptation_targets: list[str],
|
||||
) -> MetaLearningModel:
|
||||
"""Create a new meta-learning model"""
|
||||
|
||||
model_id = f"meta_{uuid4().hex[:8]}"
|
||||
|
||||
# Initialize meta-features based on adaptation targets
|
||||
meta_features = self.generate_meta_features(adaptation_targets)
|
||||
|
||||
# Set up task distributions for meta-training
|
||||
task_distributions = self.setup_task_distributions(adaptation_targets)
|
||||
|
||||
model = MetaLearningModel(
|
||||
model_id=model_id,
|
||||
model_name=model_name,
|
||||
base_algorithms=base_algorithms,
|
||||
meta_strategy=meta_strategy,
|
||||
adaptation_targets=adaptation_targets,
|
||||
meta_features=meta_features,
|
||||
task_distributions=task_distributions,
|
||||
status="training",
|
||||
)
|
||||
|
||||
session.add(model)
|
||||
session.commit()
|
||||
session.refresh(model)
|
||||
|
||||
# Start meta-training process
|
||||
asyncio.create_task(self.train_meta_model(session, model_id))
|
||||
|
||||
logger.info(f"Created meta-learning model {model_id} with strategy {meta_strategy.value}")
|
||||
return model
|
||||
|
||||
async def train_meta_model(self, session: Session, model_id: str) -> dict[str, Any]:
|
||||
"""Train a meta-learning model"""
|
||||
|
||||
model = session.execute(select(MetaLearningModel).where(MetaLearningModel.model_id == model_id)).first()
|
||||
|
||||
if not model:
|
||||
raise ValueError(f"Meta-learning model {model_id} not found")
|
||||
|
||||
try:
|
||||
# Simulate meta-training process
|
||||
training_results = await self.simulate_meta_training(model)
|
||||
|
||||
# Update model with training results
|
||||
model.meta_accuracy = training_results["accuracy"]
|
||||
model.adaptation_speed = training_results["adaptation_speed"]
|
||||
model.generalization_ability = training_results["generalization"]
|
||||
model.training_time = training_results["training_time"]
|
||||
model.computational_cost = training_results["computational_cost"]
|
||||
model.status = "ready"
|
||||
model.trained_at = datetime.now(timezone.utc)
|
||||
|
||||
session.commit()
|
||||
|
||||
logger.info(f"Meta-learning model {model_id} training completed")
|
||||
return training_results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error training meta-model {model_id}: {str(e)}")
|
||||
model.status = "failed"
|
||||
session.commit()
|
||||
raise
|
||||
|
||||
async def simulate_meta_training(self, model: MetaLearningModel) -> dict[str, Any]:
|
||||
"""Simulate meta-training process"""
|
||||
|
||||
# Simulate training time based on complexity
|
||||
base_time = 2.0 # hours
|
||||
complexity_multiplier = len(model.base_algorithms) * 0.5
|
||||
training_time = base_time * complexity_multiplier
|
||||
|
||||
# Simulate computational cost
|
||||
computational_cost = training_time * 10.0 # cost units
|
||||
|
||||
# Simulate performance metrics
|
||||
meta_accuracy = 0.75 + (len(model.adaptation_targets) * 0.05)
|
||||
adaptation_speed = 0.8 + (len(model.meta_features) * 0.02)
|
||||
generalization = 0.7 + (len(model.task_distributions) * 0.03)
|
||||
|
||||
# Cap values at 1.0
|
||||
meta_accuracy = min(1.0, meta_accuracy)
|
||||
adaptation_speed = min(1.0, adaptation_speed)
|
||||
generalization = min(1.0, generalization)
|
||||
|
||||
return {
|
||||
"accuracy": meta_accuracy,
|
||||
"adaptation_speed": adaptation_speed,
|
||||
"generalization": generalization,
|
||||
"training_time": training_time,
|
||||
"computational_cost": computational_cost,
|
||||
"convergence_epoch": int(training_time * 10),
|
||||
}
|
||||
|
||||
def generate_meta_features(self, adaptation_targets: list[str]) -> list[str]:
|
||||
"""Generate meta-features for adaptation targets"""
|
||||
|
||||
meta_features = []
|
||||
|
||||
for target in adaptation_targets:
|
||||
if target == "text_generation":
|
||||
meta_features.extend(["text_length", "complexity", "domain", "style"])
|
||||
elif target == "image_generation":
|
||||
meta_features.extend(["resolution", "style", "content_type", "complexity"])
|
||||
elif target == "reasoning":
|
||||
meta_features.extend(["logic_type", "complexity", "domain", "step_count"])
|
||||
elif target == "classification":
|
||||
meta_features.extend(["feature_count", "class_count", "data_type", "imbalance"])
|
||||
else:
|
||||
meta_features.extend(["complexity", "domain", "data_size", "quality"])
|
||||
|
||||
return list(set(meta_features))
|
||||
|
||||
def setup_task_distributions(self, adaptation_targets: list[str]) -> dict[str, float]:
|
||||
"""Set up task distributions for meta-training"""
|
||||
|
||||
distributions = {}
|
||||
total_targets = len(adaptation_targets)
|
||||
|
||||
for i, target in enumerate(adaptation_targets):
|
||||
# Distribute weights evenly with slight variations
|
||||
base_weight = 1.0 / total_targets
|
||||
variation = (i - total_targets / 2) * 0.1
|
||||
distributions[target] = max(0.1, base_weight + variation)
|
||||
|
||||
return distributions
|
||||
|
||||
async def adapt_to_new_task(
|
||||
self, session: Session, model_id: str, task_data: dict[str, Any], adaptation_steps: int = 10
|
||||
) -> dict[str, Any]:
|
||||
"""Adapt meta-learning model to new task"""
|
||||
|
||||
model = session.execute(select(MetaLearningModel).where(MetaLearningModel.model_id == model_id)).first()
|
||||
|
||||
if not model:
|
||||
raise ValueError(f"Meta-learning model {model_id} not found")
|
||||
|
||||
if model.status != "ready":
|
||||
raise ValueError(f"Model {model_id} is not ready for adaptation")
|
||||
|
||||
try:
|
||||
# Simulate adaptation process
|
||||
adaptation_results = await self.simulate_adaptation(model, task_data, adaptation_steps)
|
||||
|
||||
# Update deployment count and success rate
|
||||
model.deployment_count += 1
|
||||
model.success_rate = (
|
||||
model.success_rate * (model.deployment_count - 1) + adaptation_results["success"]
|
||||
) / model.deployment_count
|
||||
|
||||
session.commit()
|
||||
|
||||
logger.info(f"Model {model_id} adapted to new task with success rate {adaptation_results['success']:.2f}")
|
||||
return adaptation_results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error adapting model {model_id}: {str(e)}")
|
||||
raise
|
||||
|
||||
async def simulate_adaptation(self, model: MetaLearningModel, task_data: dict[str, Any], steps: int) -> dict[str, Any]:
|
||||
"""Simulate adaptation to new task"""
|
||||
|
||||
# Calculate adaptation success based on model capabilities
|
||||
base_success = model.meta_accuracy * model.adaptation_speed
|
||||
|
||||
# Factor in task similarity (simplified)
|
||||
task_similarity = 0.8 # Would calculate based on meta-features
|
||||
|
||||
# Calculate adaptation success
|
||||
adaptation_success = base_success * task_similarity * (1.0 - (0.1 / steps))
|
||||
|
||||
# Calculate adaptation time
|
||||
adaptation_time = steps * 0.1 # seconds per step
|
||||
|
||||
return {
|
||||
"success": adaptation_success,
|
||||
"adaptation_time": adaptation_time,
|
||||
"steps_used": steps,
|
||||
"final_performance": adaptation_success * 0.9, # Slight degradation
|
||||
"convergence_achieved": adaptation_success > 0.7,
|
||||
}
|
||||
|
||||
def maml_algorithm(self, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Model-Agnostic Meta-Learning algorithm"""
|
||||
|
||||
# Simplified MAML implementation
|
||||
return {
|
||||
"algorithm": "MAML",
|
||||
"inner_learning_rate": 0.01,
|
||||
"outer_learning_rate": 0.001,
|
||||
"inner_steps": 5,
|
||||
"meta_batch_size": 32,
|
||||
}
|
||||
|
||||
def reptile_algorithm(self, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Reptile algorithm implementation"""
|
||||
|
||||
return {"algorithm": "Reptile", "inner_learning_rate": 0.1, "meta_batch_size": 20, "inner_steps": 1, "epsilon": 1.0}
|
||||
|
||||
def meta_sgd_algorithm(self, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Meta-SGD algorithm implementation"""
|
||||
|
||||
return {"algorithm": "Meta-SGD", "learning_rate": 0.01, "momentum": 0.9, "weight_decay": 0.0001}
|
||||
|
||||
def prototypical_algorithm(self, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Prototypical Networks algorithm"""
|
||||
|
||||
return {
|
||||
"algorithm": "Prototypical",
|
||||
"embedding_size": 128,
|
||||
"distance_metric": "euclidean",
|
||||
"support_shots": 5,
|
||||
"query_shots": 10,
|
||||
}
|
||||
|
||||
def fast_adaptation(self, model: MetaLearningModel, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Fast adaptation strategy"""
|
||||
|
||||
return {"strategy": "fast_adaptation", "learning_rate": 0.01, "steps": 5, "adaptation_speed": 0.9}
|
||||
|
||||
def gradual_adaptation(self, model: MetaLearningModel, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Gradual adaptation strategy"""
|
||||
|
||||
return {"strategy": "gradual_adaptation", "learning_rate": 0.005, "steps": 20, "adaptation_speed": 0.7}
|
||||
|
||||
def transfer_adaptation(self, model: MetaLearningModel, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Transfer learning adaptation"""
|
||||
|
||||
return {
|
||||
"strategy": "transfer_adaptation",
|
||||
"source_tasks": model.adaptation_targets,
|
||||
"transfer_rate": 0.8,
|
||||
"fine_tuning_steps": 10,
|
||||
}
|
||||
|
||||
def multi_task_adaptation(self, model: MetaLearningModel, task_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Multi-task adaptation"""
|
||||
|
||||
return {
|
||||
"strategy": "multi_task_adaptation",
|
||||
"task_weights": model.task_distributions,
|
||||
"shared_layers": 3,
|
||||
"task_specific_layers": 2,
|
||||
}
|
||||
|
||||
|
||||
class ResourceManager:
|
||||
"""Self-optimizing resource management system"""
|
||||
|
||||
def __init__(self):
|
||||
self.optimization_algorithms = {
|
||||
"genetic_algorithm": self.genetic_optimization,
|
||||
"simulated_annealing": self.simulated_annealing,
|
||||
"gradient_descent": self.gradient_optimization,
|
||||
"bayesian_optimization": self.bayesian_optimization,
|
||||
}
|
||||
|
||||
self.resource_constraints = {
|
||||
ResourceType.CPU: {"min": 0.5, "max": 16.0, "step": 0.5},
|
||||
ResourceType.MEMORY: {"min": 1.0, "max": 64.0, "step": 1.0},
|
||||
ResourceType.GPU: {"min": 0.0, "max": 8.0, "step": 1.0},
|
||||
ResourceType.STORAGE: {"min": 10.0, "max": 1000.0, "step": 10.0},
|
||||
ResourceType.NETWORK: {"min": 10.0, "max": 1000.0, "step": 10.0},
|
||||
}
|
||||
|
||||
async def allocate_resources(
|
||||
self,
|
||||
session: Session,
|
||||
agent_id: str,
|
||||
task_requirements: dict[str, Any],
|
||||
optimization_target: OptimizationTarget = OptimizationTarget.EFFICIENCY,
|
||||
) -> ResourceAllocation:
|
||||
"""Allocate and optimize resources for agent task"""
|
||||
|
||||
allocation_id = f"alloc_{uuid4().hex[:8]}"
|
||||
|
||||
# Calculate initial resource requirements
|
||||
initial_allocation = self.calculate_initial_allocation(task_requirements)
|
||||
|
||||
# Optimize allocation based on target
|
||||
optimized_allocation = await self.optimize_allocation(initial_allocation, task_requirements, optimization_target)
|
||||
|
||||
allocation = ResourceAllocation(
|
||||
allocation_id=allocation_id,
|
||||
agent_id=agent_id,
|
||||
cpu_cores=optimized_allocation[ResourceType.CPU],
|
||||
memory_gb=optimized_allocation[ResourceType.MEMORY],
|
||||
gpu_count=optimized_allocation[ResourceType.GPU],
|
||||
gpu_memory_gb=optimized_allocation.get("gpu_memory", 0.0),
|
||||
storage_gb=optimized_allocation[ResourceType.STORAGE],
|
||||
network_bandwidth=optimized_allocation[ResourceType.NETWORK],
|
||||
optimization_target=optimization_target,
|
||||
status="allocated",
|
||||
allocated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
session.add(allocation)
|
||||
session.commit()
|
||||
session.refresh(allocation)
|
||||
|
||||
logger.info(f"Allocated resources for agent {agent_id} with target {optimization_target.value}")
|
||||
return allocation
|
||||
|
||||
def calculate_initial_allocation(self, task_requirements: dict[str, Any]) -> dict[ResourceType, float]:
|
||||
"""Calculate initial resource allocation based on task requirements"""
|
||||
|
||||
allocation = {
|
||||
ResourceType.CPU: 2.0,
|
||||
ResourceType.MEMORY: 4.0,
|
||||
ResourceType.GPU: 0.0,
|
||||
ResourceType.STORAGE: 50.0,
|
||||
ResourceType.NETWORK: 100.0,
|
||||
}
|
||||
|
||||
# Adjust based on task type
|
||||
task_type = task_requirements.get("task_type", "general")
|
||||
|
||||
if task_type == "inference":
|
||||
allocation[ResourceType.CPU] = 4.0
|
||||
allocation[ResourceType.MEMORY] = 8.0
|
||||
allocation[ResourceType.GPU] = 1.0 if task_requirements.get("model_size") == "large" else 0.0
|
||||
allocation[ResourceType.NETWORK] = 200.0
|
||||
|
||||
elif task_type == "training":
|
||||
allocation[ResourceType.CPU] = 8.0
|
||||
allocation[ResourceType.MEMORY] = 16.0
|
||||
allocation[ResourceType.GPU] = 2.0
|
||||
allocation[ResourceType.STORAGE] = 200.0
|
||||
allocation[ResourceType.NETWORK] = 500.0
|
||||
|
||||
elif task_type == "text_generation":
|
||||
allocation[ResourceType.CPU] = 2.0
|
||||
allocation[ResourceType.MEMORY] = 6.0
|
||||
allocation[ResourceType.GPU] = 0.0
|
||||
allocation[ResourceType.NETWORK] = 50.0
|
||||
|
||||
elif task_type == "image_generation":
|
||||
allocation[ResourceType.CPU] = 4.0
|
||||
allocation[ResourceType.MEMORY] = 12.0
|
||||
allocation[ResourceType.GPU] = 1.0
|
||||
allocation[ResourceType.STORAGE] = 100.0
|
||||
allocation[ResourceType.NETWORK] = 100.0
|
||||
|
||||
# Adjust based on workload size
|
||||
workload_factor = task_requirements.get("workload_factor", 1.0)
|
||||
for resource_type in allocation:
|
||||
allocation[resource_type] *= workload_factor
|
||||
|
||||
return allocation
|
||||
|
||||
async def optimize_allocation(
|
||||
self, initial_allocation: dict[ResourceType, float], task_requirements: dict[str, Any], target: OptimizationTarget
|
||||
) -> dict[ResourceType, float]:
|
||||
"""Optimize resource allocation based on target"""
|
||||
|
||||
if target == OptimizationTarget.SPEED:
|
||||
return await self.optimize_for_speed(initial_allocation, task_requirements)
|
||||
elif target == OptimizationTarget.ACCURACY:
|
||||
return await self.optimize_for_accuracy(initial_allocation, task_requirements)
|
||||
elif target == OptimizationTarget.EFFICIENCY:
|
||||
return await self.optimize_for_efficiency(initial_allocation, task_requirements)
|
||||
elif target == OptimizationTarget.COST:
|
||||
return await self.optimize_for_cost(initial_allocation, task_requirements)
|
||||
else:
|
||||
return initial_allocation
|
||||
|
||||
async def optimize_for_speed(
|
||||
self, allocation: dict[ResourceType, float], task_requirements: dict[str, Any]
|
||||
) -> dict[ResourceType, float]:
|
||||
"""Optimize allocation for speed"""
|
||||
|
||||
optimized = allocation.copy()
|
||||
|
||||
# Increase CPU and memory for faster processing
|
||||
optimized[ResourceType.CPU] = min(
|
||||
self.resource_constraints[ResourceType.CPU]["max"], optimized[ResourceType.CPU] * 1.5
|
||||
)
|
||||
optimized[ResourceType.MEMORY] = min(
|
||||
self.resource_constraints[ResourceType.MEMORY]["max"], optimized[ResourceType.MEMORY] * 1.3
|
||||
)
|
||||
|
||||
# Add GPU if available and beneficial
|
||||
if task_requirements.get("task_type") in ["inference", "image_generation"]:
|
||||
optimized[ResourceType.GPU] = min(
|
||||
self.resource_constraints[ResourceType.GPU]["max"], max(optimized[ResourceType.GPU], 1.0)
|
||||
)
|
||||
|
||||
return optimized
|
||||
|
||||
async def optimize_for_accuracy(
|
||||
self, allocation: dict[ResourceType, float], task_requirements: dict[str, Any]
|
||||
) -> dict[ResourceType, float]:
|
||||
"""Optimize allocation for accuracy"""
|
||||
|
||||
optimized = allocation.copy()
|
||||
|
||||
# Increase memory for larger models
|
||||
optimized[ResourceType.MEMORY] = min(
|
||||
self.resource_constraints[ResourceType.MEMORY]["max"], optimized[ResourceType.MEMORY] * 2.0
|
||||
)
|
||||
|
||||
# Add GPU for compute-intensive tasks
|
||||
if task_requirements.get("task_type") in ["training", "inference"]:
|
||||
optimized[ResourceType.GPU] = min(
|
||||
self.resource_constraints[ResourceType.GPU]["max"], max(optimized[ResourceType.GPU], 2.0)
|
||||
)
|
||||
optimized[ResourceType.GPU_MEMORY_GB] = optimized[ResourceType.GPU] * 8.0
|
||||
|
||||
return optimized
|
||||
|
||||
async def optimize_for_efficiency(
|
||||
self, allocation: dict[ResourceType, float], task_requirements: dict[str, Any]
|
||||
) -> dict[ResourceType, float]:
|
||||
"""Optimize allocation for efficiency"""
|
||||
|
||||
optimized = allocation.copy()
|
||||
|
||||
# Find optimal balance between resources
|
||||
task_type = task_requirements.get("task_type", "general")
|
||||
|
||||
if task_type == "text_generation":
|
||||
# Text generation is CPU-efficient
|
||||
optimized[ResourceType.CPU] = max(
|
||||
self.resource_constraints[ResourceType.CPU]["min"], optimized[ResourceType.CPU] * 0.8
|
||||
)
|
||||
optimized[ResourceType.GPU] = 0.0
|
||||
|
||||
elif task_type == "inference":
|
||||
# Moderate GPU usage for inference
|
||||
optimized[ResourceType.GPU] = min(
|
||||
self.resource_constraints[ResourceType.GPU]["max"], max(0.5, optimized[ResourceType.GPU] * 0.7)
|
||||
)
|
||||
|
||||
return optimized
|
||||
|
||||
async def optimize_for_cost(
|
||||
self, allocation: dict[ResourceType, float], task_requirements: dict[str, Any]
|
||||
) -> dict[ResourceType, float]:
|
||||
"""Optimize allocation for cost"""
|
||||
|
||||
optimized = allocation.copy()
|
||||
|
||||
# Minimize expensive resources
|
||||
optimized[ResourceType.GPU] = 0.0
|
||||
optimized[ResourceType.CPU] = max(
|
||||
self.resource_constraints[ResourceType.CPU]["min"], optimized[ResourceType.CPU] * 0.5
|
||||
)
|
||||
optimized[ResourceType.MEMORY] = max(
|
||||
self.resource_constraints[ResourceType.MEMORY]["min"], optimized[ResourceType.MEMORY] * 0.7
|
||||
)
|
||||
|
||||
return optimized
|
||||
|
||||
def genetic_optimization(self, allocation: dict[ResourceType, float]) -> dict[str, Any]:
|
||||
"""Genetic algorithm for resource optimization"""
|
||||
|
||||
return {
|
||||
"algorithm": "genetic_algorithm",
|
||||
"population_size": 50,
|
||||
"generations": 100,
|
||||
"mutation_rate": 0.1,
|
||||
"crossover_rate": 0.8,
|
||||
}
|
||||
|
||||
def simulated_annealing(self, allocation: dict[ResourceType, float]) -> dict[str, Any]:
|
||||
"""Simulated annealing optimization"""
|
||||
|
||||
return {"algorithm": "simulated_annealing", "initial_temperature": 100.0, "cooling_rate": 0.95, "iterations": 1000}
|
||||
|
||||
def gradient_optimization(self, allocation: dict[ResourceType, float]) -> dict[str, Any]:
|
||||
"""Gradient descent optimization"""
|
||||
|
||||
return {"algorithm": "gradient_descent", "learning_rate": 0.01, "iterations": 500, "momentum": 0.9}
|
||||
|
||||
def bayesian_optimization(self, allocation: dict[ResourceType, float]) -> dict[str, Any]:
|
||||
"""Bayesian optimization"""
|
||||
|
||||
return {
|
||||
"algorithm": "bayesian_optimization",
|
||||
"acquisition_function": "expected_improvement",
|
||||
"iterations": 50,
|
||||
"exploration_weight": 0.1,
|
||||
}
|
||||
|
||||
|
||||
class PerformanceOptimizer:
|
||||
"""Advanced performance optimization system"""
|
||||
|
||||
def __init__(self):
|
||||
self.optimization_techniques = {
|
||||
"hyperparameter_tuning": self.tune_hyperparameters,
|
||||
"architecture_optimization": self.optimize_architecture,
|
||||
"algorithm_selection": self.select_algorithm,
|
||||
"data_optimization": self.optimize_data_pipeline,
|
||||
}
|
||||
|
||||
self.performance_targets = {
|
||||
PerformanceMetric.ACCURACY: {"weight": 0.3, "target": 0.95},
|
||||
PerformanceMetric.LATENCY: {"weight": 0.25, "target": 100.0}, # ms
|
||||
PerformanceMetric.THROUGHPUT: {"weight": 0.2, "target": 100.0},
|
||||
PerformanceMetric.RESOURCE_EFFICIENCY: {"weight": 0.15, "target": 0.8},
|
||||
PerformanceMetric.COST_EFFICIENCY: {"weight": 0.1, "target": 0.9},
|
||||
}
|
||||
|
||||
async def optimize_agent_performance(
|
||||
self, session: Session, agent_id: str, target_metric: PerformanceMetric, current_performance: dict[str, float]
|
||||
) -> PerformanceOptimization:
|
||||
"""Optimize agent performance for specific metric"""
|
||||
|
||||
optimization_id = f"opt_{uuid4().hex[:8]}"
|
||||
|
||||
# Create optimization record
|
||||
optimization = PerformanceOptimization(
|
||||
optimization_id=optimization_id,
|
||||
agent_id=agent_id,
|
||||
optimization_type="comprehensive",
|
||||
target_metric=target_metric,
|
||||
baseline_performance=current_performance,
|
||||
baseline_cost=self.calculate_cost(current_performance),
|
||||
status="running",
|
||||
)
|
||||
|
||||
session.add(optimization)
|
||||
session.commit()
|
||||
session.refresh(optimization)
|
||||
|
||||
try:
|
||||
# Run optimization process
|
||||
optimization_results = await self.run_optimization_process(agent_id, target_metric, current_performance)
|
||||
|
||||
# Update optimization with results
|
||||
optimization.optimized_performance = optimization_results["performance"]
|
||||
optimization.optimized_resources = optimization_results["resources"]
|
||||
optimization.optimized_cost = optimization_results["cost"]
|
||||
optimization.performance_improvement = optimization_results["improvement"]
|
||||
optimization.resource_savings = optimization_results["savings"]
|
||||
optimization.cost_savings = optimization_results["cost_savings"]
|
||||
optimization.overall_efficiency_gain = optimization_results["efficiency_gain"]
|
||||
optimization.optimization_duration = optimization_results["duration"]
|
||||
optimization.iterations_required = optimization_results["iterations"]
|
||||
optimization.convergence_achieved = optimization_results["converged"]
|
||||
optimization.optimization_applied = True
|
||||
optimization.status = "completed"
|
||||
optimization.completed_at = datetime.now(timezone.utc)
|
||||
|
||||
session.commit()
|
||||
|
||||
logger.info(f"Performance optimization {optimization_id} completed for agent {agent_id}")
|
||||
return optimization
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error optimizing performance for agent {agent_id}: {str(e)}")
|
||||
optimization.status = "failed"
|
||||
session.commit()
|
||||
raise
|
||||
|
||||
async def run_optimization_process(
|
||||
self, agent_id: str, target_metric: PerformanceMetric, current_performance: dict[str, float]
|
||||
) -> dict[str, Any]:
|
||||
"""Run comprehensive optimization process"""
|
||||
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
# Step 1: Analyze current performance
|
||||
analysis_results = self.analyze_current_performance(current_performance, target_metric)
|
||||
|
||||
# Step 2: Generate optimization candidates
|
||||
candidates = await self.generate_optimization_candidates(target_metric, analysis_results)
|
||||
|
||||
# Step 3: Evaluate candidates
|
||||
best_candidate = await self.evaluate_candidates(candidates, target_metric)
|
||||
|
||||
# Step 4: Apply optimization
|
||||
applied_performance = await self.apply_optimization(best_candidate)
|
||||
|
||||
# Step 5: Calculate improvements
|
||||
improvements = self.calculate_improvements(current_performance, applied_performance)
|
||||
|
||||
end_time = datetime.now(timezone.utc)
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"performance": applied_performance,
|
||||
"resources": best_candidate.get("resources", {}),
|
||||
"cost": self.calculate_cost(applied_performance),
|
||||
"improvement": improvements["overall"],
|
||||
"savings": improvements["resource"],
|
||||
"cost_savings": improvements["cost"],
|
||||
"efficiency_gain": improvements["efficiency"],
|
||||
"duration": duration,
|
||||
"iterations": len(candidates),
|
||||
"converged": improvements["overall"] > 0.05,
|
||||
}
|
||||
|
||||
def analyze_current_performance(
|
||||
self, current_performance: dict[str, float], target_metric: PerformanceMetric
|
||||
) -> dict[str, Any]:
|
||||
"""Analyze current performance to identify bottlenecks"""
|
||||
|
||||
analysis = {
|
||||
"current_value": current_performance.get(target_metric.value, 0.0),
|
||||
"target_value": self.performance_targets[target_metric]["target"],
|
||||
"gap": 0.0,
|
||||
"bottlenecks": [],
|
||||
"improvement_potential": 0.0,
|
||||
}
|
||||
|
||||
# Calculate performance gap
|
||||
current_value = analysis["current_value"]
|
||||
target_value = analysis["target_value"]
|
||||
|
||||
if target_metric == PerformanceMetric.ACCURACY:
|
||||
analysis["gap"] = target_value - current_value
|
||||
analysis["improvement_potential"] = min(1.0, analysis["gap"] / target_value)
|
||||
elif target_metric == PerformanceMetric.LATENCY:
|
||||
analysis["gap"] = current_value - target_value
|
||||
analysis["improvement_potential"] = min(1.0, analysis["gap"] / current_value)
|
||||
else:
|
||||
# For other metrics, calculate relative improvement
|
||||
analysis["gap"] = target_value - current_value
|
||||
analysis["improvement_potential"] = min(1.0, analysis["gap"] / target_value)
|
||||
|
||||
# Identify bottlenecks
|
||||
if current_performance.get("cpu_utilization", 0) > 0.9:
|
||||
analysis["bottlenecks"].append("cpu")
|
||||
if current_performance.get("memory_utilization", 0) > 0.9:
|
||||
analysis["bottlenecks"].append("memory")
|
||||
if current_performance.get("gpu_utilization", 0) > 0.9:
|
||||
analysis["bottlenecks"].append("gpu")
|
||||
|
||||
return analysis
|
||||
|
||||
async def generate_optimization_candidates(
|
||||
self, target_metric: PerformanceMetric, analysis: dict[str, Any]
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Generate optimization candidates"""
|
||||
|
||||
candidates = []
|
||||
|
||||
# Hyperparameter tuning candidate
|
||||
hp_candidate = await self.tune_hyperparameters(target_metric, analysis)
|
||||
candidates.append(hp_candidate)
|
||||
|
||||
# Architecture optimization candidate
|
||||
arch_candidate = await self.optimize_architecture(target_metric, analysis)
|
||||
candidates.append(arch_candidate)
|
||||
|
||||
# Algorithm selection candidate
|
||||
algo_candidate = await self.select_algorithm(target_metric, analysis)
|
||||
candidates.append(algo_candidate)
|
||||
|
||||
# Data optimization candidate
|
||||
data_candidate = await self.optimize_data_pipeline(target_metric, analysis)
|
||||
candidates.append(data_candidate)
|
||||
|
||||
return candidates
|
||||
|
||||
async def evaluate_candidates(self, candidates: list[dict[str, Any]], target_metric: PerformanceMetric) -> dict[str, Any]:
|
||||
"""Evaluate optimization candidates and select best"""
|
||||
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
|
||||
for candidate in candidates:
|
||||
# Calculate expected performance improvement
|
||||
expected_improvement = candidate.get("expected_improvement", 0.0)
|
||||
resource_cost = candidate.get("resource_cost", 1.0)
|
||||
implementation_complexity = candidate.get("complexity", 0.5)
|
||||
|
||||
# Calculate overall score
|
||||
score = expected_improvement * 0.6 - resource_cost * 0.2 - implementation_complexity * 0.2
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_candidate = candidate
|
||||
|
||||
return best_candidate or {}
|
||||
|
||||
async def apply_optimization(self, candidate: dict[str, Any]) -> dict[str, float]:
|
||||
"""Apply optimization and return expected performance"""
|
||||
|
||||
# Simulate applying optimization
|
||||
base_performance = candidate.get("base_performance", {})
|
||||
improvement_factor = candidate.get("expected_improvement", 0.0)
|
||||
|
||||
applied_performance = {}
|
||||
for metric, value in base_performance.items():
|
||||
if metric == candidate.get("target_metric"):
|
||||
applied_performance[metric] = value * (1.0 + improvement_factor)
|
||||
else:
|
||||
# Other metrics may change slightly
|
||||
applied_performance[metric] = value * (1.0 + improvement_factor * 0.1)
|
||||
|
||||
return applied_performance
|
||||
|
||||
def calculate_improvements(self, baseline: dict[str, float], optimized: dict[str, float]) -> dict[str, float]:
|
||||
"""Calculate performance improvements"""
|
||||
|
||||
improvements = {"overall": 0.0, "resource": 0.0, "cost": 0.0, "efficiency": 0.0}
|
||||
|
||||
# Calculate overall improvement
|
||||
baseline_total = sum(baseline.values())
|
||||
optimized_total = sum(optimized.values())
|
||||
improvements["overall"] = (optimized_total - baseline_total) / baseline_total if baseline_total > 0 else 0.0
|
||||
|
||||
# Calculate resource savings (simplified)
|
||||
baseline_resources = baseline.get("cpu_cores", 1.0) + baseline.get("memory_gb", 2.0)
|
||||
optimized_resources = optimized.get("cpu_cores", 1.0) + optimized.get("memory_gb", 2.0)
|
||||
improvements["resource"] = (
|
||||
(baseline_resources - optimized_resources) / baseline_resources if baseline_resources > 0 else 0.0
|
||||
)
|
||||
|
||||
# Calculate cost savings
|
||||
baseline_cost = self.calculate_cost(baseline)
|
||||
optimized_cost = self.calculate_cost(optimized)
|
||||
improvements["cost"] = (baseline_cost - optimized_cost) / baseline_cost if baseline_cost > 0 else 0.0
|
||||
|
||||
# Calculate efficiency gain
|
||||
improvements["efficiency"] = improvements["overall"] + improvements["resource"] + improvements["cost"]
|
||||
|
||||
return improvements
|
||||
|
||||
def calculate_cost(self, performance: dict[str, float]) -> float:
|
||||
"""Calculate cost based on resource usage"""
|
||||
|
||||
cpu_cost = performance.get("cpu_cores", 1.0) * 10.0 # $10 per core
|
||||
memory_cost = performance.get("memory_gb", 2.0) * 2.0 # $2 per GB
|
||||
gpu_cost = performance.get("gpu_count", 0.0) * 100.0 # $100 per GPU
|
||||
storage_cost = performance.get("storage_gb", 50.0) * 0.1 # $0.1 per GB
|
||||
|
||||
return cpu_cost + memory_cost + gpu_cost + storage_cost
|
||||
|
||||
async def tune_hyperparameters(self, target_metric: PerformanceMetric, analysis: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Tune hyperparameters for performance optimization"""
|
||||
|
||||
return {
|
||||
"technique": "hyperparameter_tuning",
|
||||
"target_metric": target_metric.value,
|
||||
"parameters": {"learning_rate": 0.001, "batch_size": 64, "dropout_rate": 0.1, "weight_decay": 0.0001},
|
||||
"expected_improvement": 0.15,
|
||||
"resource_cost": 0.1,
|
||||
"complexity": 0.3,
|
||||
}
|
||||
|
||||
async def optimize_architecture(self, target_metric: PerformanceMetric, analysis: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Optimize model architecture"""
|
||||
|
||||
return {
|
||||
"technique": "architecture_optimization",
|
||||
"target_metric": target_metric.value,
|
||||
"architecture": {"layers": [256, 128, 64], "activations": ["relu", "relu", "tanh"], "normalization": "batch_norm"},
|
||||
"expected_improvement": 0.25,
|
||||
"resource_cost": 0.2,
|
||||
"complexity": 0.7,
|
||||
}
|
||||
|
||||
async def select_algorithm(self, target_metric: PerformanceMetric, analysis: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Select optimal algorithm"""
|
||||
|
||||
return {
|
||||
"technique": "algorithm_selection",
|
||||
"target_metric": target_metric.value,
|
||||
"algorithm": "transformer",
|
||||
"expected_improvement": 0.20,
|
||||
"resource_cost": 0.3,
|
||||
"complexity": 0.5,
|
||||
}
|
||||
|
||||
async def optimize_data_pipeline(self, target_metric: PerformanceMetric, analysis: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Optimize data processing pipeline"""
|
||||
|
||||
return {
|
||||
"technique": "data_optimization",
|
||||
"target_metric": target_metric.value,
|
||||
"optimizations": {"data_augmentation": True, "batch_normalization": True, "early_stopping": True},
|
||||
"expected_improvement": 0.10,
|
||||
"resource_cost": 0.05,
|
||||
"complexity": 0.2,
|
||||
}
|
||||
|
||||
|
||||
class AgentPerformanceService:
|
||||
"""Main service for advanced agent performance management"""
|
||||
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
self.meta_learning_engine = MetaLearningEngine()
|
||||
self.resource_manager = ResourceManager()
|
||||
self.performance_optimizer = PerformanceOptimizer()
|
||||
|
||||
async def create_performance_profile(
|
||||
self, agent_id: str, agent_type: str = "hermes", initial_metrics: dict[str, float] | None = None
|
||||
) -> AgentPerformanceProfile:
|
||||
"""Create comprehensive agent performance profile"""
|
||||
|
||||
profile_id = f"perf_{uuid4().hex[:8]}"
|
||||
|
||||
profile = AgentPerformanceProfile(
|
||||
profile_id=profile_id,
|
||||
agent_id=agent_id,
|
||||
agent_type=agent_type,
|
||||
performance_metrics=initial_metrics or {},
|
||||
learning_strategies=["meta_learning", "transfer_learning"],
|
||||
specialization_areas=["general"],
|
||||
expertise_levels={},
|
||||
performance_history=[],
|
||||
benchmark_scores={},
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
self.session.add(profile)
|
||||
self.session.commit()
|
||||
self.session.refresh(profile)
|
||||
|
||||
logger.info(f"Created performance profile {profile_id} for agent {agent_id}")
|
||||
return profile
|
||||
|
||||
async def update_performance_metrics(
|
||||
self, agent_id: str, new_metrics: dict[str, float], task_context: dict[str, Any] | None = None
|
||||
) -> AgentPerformanceProfile:
|
||||
"""Update agent performance metrics"""
|
||||
|
||||
profile = self.session.execute(
|
||||
select(AgentPerformanceProfile).where(AgentPerformanceProfile.agent_id == agent_id)
|
||||
).first()
|
||||
|
||||
if not profile:
|
||||
# Create profile if it doesn't exist
|
||||
profile = await self.create_performance_profile(agent_id, "hermes", new_metrics)
|
||||
else:
|
||||
# Update existing profile
|
||||
profile.performance_metrics.update(new_metrics)
|
||||
|
||||
# Add to performance history
|
||||
history_entry = {"timestamp": datetime.now(timezone.utc).isoformat(), "metrics": new_metrics, "context": task_context or {}}
|
||||
profile.performance_history.append(history_entry)
|
||||
|
||||
# Calculate overall score
|
||||
profile.overall_score = self.calculate_overall_score(profile.performance_metrics)
|
||||
|
||||
# Update trends
|
||||
profile.improvement_trends = self.calculate_improvement_trends(profile.performance_history)
|
||||
|
||||
profile.updated_at = datetime.now(timezone.utc)
|
||||
profile.last_assessed = datetime.now(timezone.utc)
|
||||
|
||||
self.session.commit()
|
||||
|
||||
return profile
|
||||
|
||||
def calculate_overall_score(self, metrics: dict[str, float]) -> float:
|
||||
"""Calculate overall performance score"""
|
||||
|
||||
if not metrics:
|
||||
return 0.0
|
||||
|
||||
# Weight different metrics
|
||||
weights = {
|
||||
"accuracy": 0.3,
|
||||
"latency": -0.2, # Lower is better
|
||||
"throughput": 0.2,
|
||||
"efficiency": 0.15,
|
||||
"cost_efficiency": 0.15,
|
||||
}
|
||||
|
||||
score = 0.0
|
||||
total_weight = 0.0
|
||||
|
||||
for metric, value in metrics.items():
|
||||
weight = weights.get(metric, 0.1)
|
||||
score += value * weight
|
||||
total_weight += weight
|
||||
|
||||
return score / total_weight if total_weight > 0 else 0.0
|
||||
|
||||
def calculate_improvement_trends(self, history: list[dict[str, Any]]) -> dict[str, float]:
|
||||
"""Calculate performance improvement trends"""
|
||||
|
||||
if len(history) < 2:
|
||||
return {}
|
||||
|
||||
trends = {}
|
||||
|
||||
# Get latest and previous metrics
|
||||
latest_metrics = history[-1]["metrics"]
|
||||
previous_metrics = history[-2]["metrics"]
|
||||
|
||||
for metric in latest_metrics:
|
||||
if metric in previous_metrics:
|
||||
latest_value = latest_metrics[metric]
|
||||
previous_value = previous_metrics[metric]
|
||||
|
||||
if previous_value != 0:
|
||||
change = (latest_value - previous_value) / abs(previous_value)
|
||||
trends[metric] = change
|
||||
|
||||
return trends
|
||||
|
||||
async def get_comprehensive_profile(self, agent_id: str) -> dict[str, Any]:
|
||||
"""Get comprehensive agent performance profile"""
|
||||
|
||||
profile = self.session.execute(
|
||||
select(AgentPerformanceProfile).where(AgentPerformanceProfile.agent_id == agent_id)
|
||||
).first()
|
||||
|
||||
if not profile:
|
||||
return {"error": "Profile not found"}
|
||||
|
||||
return {
|
||||
"profile_id": profile.profile_id,
|
||||
"agent_id": profile.agent_id,
|
||||
"agent_type": profile.agent_type,
|
||||
"overall_score": profile.overall_score,
|
||||
"performance_metrics": profile.performance_metrics,
|
||||
"learning_strategies": profile.learning_strategies,
|
||||
"specialization_areas": profile.specialization_areas,
|
||||
"expertise_levels": profile.expertise_levels,
|
||||
"resource_efficiency": profile.resource_efficiency,
|
||||
"cost_per_task": profile.cost_per_task,
|
||||
"throughput": profile.throughput,
|
||||
"average_latency": profile.average_latency,
|
||||
"performance_history": profile.performance_history,
|
||||
"improvement_trends": profile.improvement_trends,
|
||||
"benchmark_scores": profile.benchmark_scores,
|
||||
"ranking_position": profile.ranking_position,
|
||||
"percentile_rank": profile.percentile_rank,
|
||||
"last_assessed": profile.last_assessed.isoformat() if profile.last_assessed else None,
|
||||
}
|
||||
560
apps/agent-management/src/app/services/agent_portfolio_manager.py
Executable file
560
apps/agent-management/src/app/services/agent_portfolio_manager.py
Executable file
@@ -0,0 +1,560 @@
|
||||
"""
|
||||
Agent Portfolio Manager Service
|
||||
|
||||
Advanced portfolio management for autonomous AI agents in the AITBC ecosystem.
|
||||
Provides portfolio creation, rebalancing, risk assessment, and trading strategy execution.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
from aitbc import get_logger
|
||||
from fastapi import HTTPException
|
||||
from sqlalchemy import select
|
||||
from sqlmodel import Session
|
||||
|
||||
from ..blockchain.contract_interactions import ContractInteractionService
|
||||
from app.domain.agent_portfolio import (
|
||||
AgentPortfolio,
|
||||
PortfolioAsset,
|
||||
PortfolioStrategy,
|
||||
PortfolioTrade,
|
||||
RiskMetrics,
|
||||
TradeStatus,
|
||||
)
|
||||
from ..marketdata.price_service import PriceService
|
||||
from ..ml.strategy_optimizer import StrategyOptimizer
|
||||
from ..risk.risk_calculator import RiskCalculator
|
||||
from ..schemas.portfolio import (
|
||||
PortfolioCreate,
|
||||
PortfolioResponse,
|
||||
RebalanceRequest,
|
||||
RebalanceResponse,
|
||||
RiskAssessmentResponse,
|
||||
StrategyCreate,
|
||||
StrategyResponse,
|
||||
TradeRequest,
|
||||
TradeResponse,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AgentPortfolioManager:
|
||||
"""Advanced portfolio management for autonomous agents"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session: Session,
|
||||
contract_service: ContractInteractionService,
|
||||
price_service: PriceService,
|
||||
risk_calculator: RiskCalculator,
|
||||
strategy_optimizer: StrategyOptimizer,
|
||||
) -> None:
|
||||
self.session = session
|
||||
self.contract_service = contract_service
|
||||
self.price_service = price_service
|
||||
self.risk_calculator = risk_calculator
|
||||
self.strategy_optimizer = strategy_optimizer
|
||||
|
||||
async def create_portfolio(self, portfolio_data: PortfolioCreate, agent_address: str) -> PortfolioResponse:
|
||||
"""Create a new portfolio for an autonomous agent"""
|
||||
|
||||
try:
|
||||
# Validate agent address
|
||||
if not self._is_valid_address(agent_address):
|
||||
raise HTTPException(status_code=400, detail="Invalid agent address")
|
||||
|
||||
# Check if portfolio already exists
|
||||
existing_portfolio = self.session.execute(
|
||||
select(AgentPortfolio).where(AgentPortfolio.agent_address == agent_address)
|
||||
).first()
|
||||
|
||||
if existing_portfolio:
|
||||
raise HTTPException(status_code=400, detail="Portfolio already exists for this agent")
|
||||
|
||||
# Get strategy
|
||||
strategy = self.session.get(PortfolioStrategy, portfolio_data.strategy_id)
|
||||
if not strategy or not strategy.is_active:
|
||||
raise HTTPException(status_code=404, detail="Strategy not found")
|
||||
|
||||
# Create portfolio
|
||||
portfolio = AgentPortfolio(
|
||||
agent_address=agent_address,
|
||||
strategy_id=portfolio_data.strategy_id,
|
||||
initial_capital=portfolio_data.initial_capital,
|
||||
risk_tolerance=portfolio_data.risk_tolerance,
|
||||
is_active=True,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
last_rebalance=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
self.session.add(portfolio)
|
||||
self.session.commit()
|
||||
self.session.refresh(portfolio)
|
||||
|
||||
# Initialize portfolio assets based on strategy
|
||||
await self._initialize_portfolio_assets(portfolio, strategy)
|
||||
|
||||
# Deploy smart contract portfolio
|
||||
contract_portfolio_id = await self._deploy_contract_portfolio(portfolio, agent_address, strategy)
|
||||
|
||||
portfolio.contract_portfolio_id = contract_portfolio_id
|
||||
self.session.commit()
|
||||
|
||||
logger.info(f"Created portfolio {portfolio.id} for agent {agent_address}")
|
||||
|
||||
return PortfolioResponse.from_orm(portfolio)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating portfolio: {str(e)}")
|
||||
self.session.rollback()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
async def execute_trade(self, trade_request: TradeRequest, agent_address: str) -> TradeResponse:
|
||||
"""Execute a trade within the agent's portfolio"""
|
||||
|
||||
try:
|
||||
# Get portfolio
|
||||
portfolio = self._get_agent_portfolio(agent_address)
|
||||
|
||||
# Validate trade request
|
||||
validation_result = await self._validate_trade_request(portfolio, trade_request)
|
||||
if not validation_result.is_valid:
|
||||
raise HTTPException(status_code=400, detail=validation_result.error_message)
|
||||
|
||||
# Get current prices
|
||||
sell_price = await self.price_service.get_price(trade_request.sell_token)
|
||||
buy_price = await self.price_service.get_price(trade_request.buy_token)
|
||||
|
||||
# Calculate expected buy amount
|
||||
expected_buy_amount = self._calculate_buy_amount(trade_request.sell_amount, sell_price, buy_price)
|
||||
|
||||
# Check slippage
|
||||
if expected_buy_amount < trade_request.min_buy_amount:
|
||||
raise HTTPException(status_code=400, detail="Insufficient buy amount (slippage protection)")
|
||||
|
||||
# Execute trade on blockchain
|
||||
trade_result = await self.contract_service.execute_portfolio_trade(
|
||||
portfolio.contract_portfolio_id,
|
||||
trade_request.sell_token,
|
||||
trade_request.buy_token,
|
||||
trade_request.sell_amount,
|
||||
trade_request.min_buy_amount,
|
||||
)
|
||||
|
||||
# Record trade in database
|
||||
trade = PortfolioTrade(
|
||||
portfolio_id=portfolio.id,
|
||||
sell_token=trade_request.sell_token,
|
||||
buy_token=trade_request.buy_token,
|
||||
sell_amount=trade_request.sell_amount,
|
||||
buy_amount=trade_result.buy_amount,
|
||||
price=trade_result.price,
|
||||
status=TradeStatus.EXECUTED,
|
||||
transaction_hash=trade_result.transaction_hash,
|
||||
executed_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
self.session.add(trade)
|
||||
|
||||
# Update portfolio assets
|
||||
await self._update_portfolio_assets(portfolio, trade)
|
||||
|
||||
# Update portfolio value and risk
|
||||
await self._update_portfolio_metrics(portfolio)
|
||||
|
||||
self.session.commit()
|
||||
self.session.refresh(trade)
|
||||
|
||||
logger.info(f"Executed trade {trade.id} for portfolio {portfolio.id}")
|
||||
|
||||
return TradeResponse.from_orm(trade)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing trade: {str(e)}")
|
||||
self.session.rollback()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
async def execute_rebalancing(self, rebalance_request: RebalanceRequest, agent_address: str) -> RebalanceResponse:
|
||||
"""Automated portfolio rebalancing based on market conditions"""
|
||||
|
||||
try:
|
||||
# Get portfolio
|
||||
portfolio = self._get_agent_portfolio(agent_address)
|
||||
|
||||
# Check if rebalancing is needed
|
||||
if not await self._needs_rebalancing(portfolio):
|
||||
return RebalanceResponse(success=False, message="Rebalancing not needed at this time")
|
||||
|
||||
# Get current market conditions
|
||||
market_conditions = await self.price_service.get_market_conditions()
|
||||
|
||||
# Calculate optimal allocations
|
||||
optimal_allocations = await self.strategy_optimizer.calculate_optimal_allocations(portfolio, market_conditions)
|
||||
|
||||
# Generate rebalancing trades
|
||||
rebalance_trades = await self._generate_rebalance_trades(portfolio, optimal_allocations)
|
||||
|
||||
if not rebalance_trades:
|
||||
return RebalanceResponse(success=False, message="No rebalancing trades required")
|
||||
|
||||
# Execute rebalancing trades
|
||||
executed_trades = []
|
||||
for trade in rebalance_trades:
|
||||
try:
|
||||
trade_response = await self.execute_trade(trade, agent_address)
|
||||
executed_trades.append(trade_response)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to execute rebalancing trade: {str(e)}")
|
||||
continue
|
||||
|
||||
# Update portfolio rebalance timestamp
|
||||
portfolio.last_rebalance = datetime.now(timezone.utc)
|
||||
self.session.commit()
|
||||
|
||||
logger.info(f"Rebalanced portfolio {portfolio.id} with {len(executed_trades)} trades")
|
||||
|
||||
return RebalanceResponse(
|
||||
success=True, message=f"Rebalanced with {len(executed_trades)} trades", trades_executed=len(executed_trades)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing rebalancing: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
async def risk_assessment(self, agent_address: str) -> RiskAssessmentResponse:
|
||||
"""Real-time risk assessment and position sizing"""
|
||||
|
||||
try:
|
||||
# Get portfolio
|
||||
portfolio = self._get_agent_portfolio(agent_address)
|
||||
|
||||
# Get current portfolio value
|
||||
portfolio_value = await self._calculate_portfolio_value(portfolio)
|
||||
|
||||
# Calculate risk metrics
|
||||
risk_metrics = await self.risk_calculator.calculate_portfolio_risk(portfolio, portfolio_value)
|
||||
|
||||
# Update risk metrics in database
|
||||
existing_metrics = self.session.execute(
|
||||
select(RiskMetrics).where(RiskMetrics.portfolio_id == portfolio.id)
|
||||
).first()
|
||||
|
||||
if existing_metrics:
|
||||
existing_metrics.volatility = risk_metrics.volatility
|
||||
existing_metrics.max_drawdown = risk_metrics.max_drawdown
|
||||
existing_metrics.sharpe_ratio = risk_metrics.sharpe_ratio
|
||||
existing_metrics.var_95 = risk_metrics.var_95
|
||||
existing_metrics.risk_level = risk_metrics.risk_level
|
||||
existing_metrics.updated_at = datetime.now(timezone.utc)
|
||||
else:
|
||||
risk_metrics.portfolio_id = portfolio.id
|
||||
risk_metrics.updated_at = datetime.now(timezone.utc)
|
||||
self.session.add(risk_metrics)
|
||||
|
||||
# Update portfolio risk score
|
||||
portfolio.risk_score = risk_metrics.overall_risk_score
|
||||
self.session.commit()
|
||||
|
||||
logger.info(f"Risk assessment completed for portfolio {portfolio.id}")
|
||||
|
||||
return RiskAssessmentResponse.from_orm(risk_metrics)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in risk assessment: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
async def get_portfolio_performance(self, agent_address: str, period: str = "30d") -> dict:
|
||||
"""Get portfolio performance metrics"""
|
||||
|
||||
try:
|
||||
# Get portfolio
|
||||
portfolio = self._get_agent_portfolio(agent_address)
|
||||
|
||||
# Calculate performance metrics
|
||||
performance_data = await self._calculate_performance_metrics(portfolio, period)
|
||||
|
||||
return performance_data
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting portfolio performance: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
async def create_portfolio_strategy(self, strategy_data: StrategyCreate) -> StrategyResponse:
|
||||
"""Create a new portfolio strategy"""
|
||||
|
||||
try:
|
||||
# Validate strategy allocations
|
||||
total_allocation = sum(strategy_data.target_allocations.values())
|
||||
if abs(total_allocation - 100.0) > 0.01: # Allow small rounding errors
|
||||
raise HTTPException(status_code=400, detail="Target allocations must sum to 100%")
|
||||
|
||||
# Create strategy
|
||||
strategy = PortfolioStrategy(
|
||||
name=strategy_data.name,
|
||||
strategy_type=strategy_data.strategy_type,
|
||||
target_allocations=strategy_data.target_allocations,
|
||||
max_drawdown=strategy_data.max_drawdown,
|
||||
rebalance_frequency=strategy_data.rebalance_frequency,
|
||||
is_active=True,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
self.session.add(strategy)
|
||||
self.session.commit()
|
||||
self.session.refresh(strategy)
|
||||
|
||||
logger.info(f"Created strategy {strategy.id}: {strategy.name}")
|
||||
|
||||
return StrategyResponse.from_orm(strategy)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating strategy: {str(e)}")
|
||||
self.session.rollback()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# Private helper methods
|
||||
|
||||
def _get_agent_portfolio(self, agent_address: str) -> AgentPortfolio:
|
||||
"""Get portfolio for agent address"""
|
||||
portfolio = self.session.execute(select(AgentPortfolio).where(AgentPortfolio.agent_address == agent_address)).first()
|
||||
|
||||
if not portfolio:
|
||||
raise HTTPException(status_code=404, detail="Portfolio not found")
|
||||
|
||||
return portfolio
|
||||
|
||||
def _is_valid_address(self, address: str) -> bool:
|
||||
"""Validate Ethereum address"""
|
||||
return address.startswith("0x") and len(address) == 42 and all(c in "0123456789abcdefABCDEF" for c in address[2:])
|
||||
|
||||
async def _initialize_portfolio_assets(self, portfolio: AgentPortfolio, strategy: PortfolioStrategy) -> None:
|
||||
"""Initialize portfolio assets based on strategy allocations"""
|
||||
|
||||
for token_symbol, allocation in strategy.target_allocations.items():
|
||||
if allocation > 0:
|
||||
asset = PortfolioAsset(
|
||||
portfolio_id=portfolio.id,
|
||||
token_symbol=token_symbol,
|
||||
target_allocation=allocation,
|
||||
current_allocation=0.0,
|
||||
balance=0,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
self.session.add(asset)
|
||||
|
||||
async def _deploy_contract_portfolio(
|
||||
self, portfolio: AgentPortfolio, agent_address: str, strategy: PortfolioStrategy
|
||||
) -> str:
|
||||
"""Deploy smart contract portfolio"""
|
||||
|
||||
try:
|
||||
# Convert strategy allocations to contract format
|
||||
contract_allocations = {
|
||||
token: int(allocation * 100) # Convert to basis points
|
||||
for token, allocation in strategy.target_allocations.items()
|
||||
}
|
||||
|
||||
# Create portfolio on blockchain
|
||||
portfolio_id = await self.contract_service.create_portfolio(
|
||||
agent_address, strategy.strategy_type.value, contract_allocations
|
||||
)
|
||||
|
||||
return str(portfolio_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error deploying contract portfolio: {str(e)}")
|
||||
raise
|
||||
|
||||
async def _validate_trade_request(self, portfolio: AgentPortfolio, trade_request: TradeRequest) -> ValidationResult:
|
||||
"""Validate trade request"""
|
||||
|
||||
# Check if sell token exists in portfolio
|
||||
sell_asset = self.session.execute(
|
||||
select(PortfolioAsset).where(
|
||||
PortfolioAsset.portfolio_id == portfolio.id, PortfolioAsset.token_symbol == trade_request.sell_token
|
||||
)
|
||||
).first()
|
||||
|
||||
if not sell_asset:
|
||||
return ValidationResult(is_valid=False, error_message="Sell token not found in portfolio")
|
||||
|
||||
# Check sufficient balance
|
||||
if sell_asset.balance < trade_request.sell_amount:
|
||||
return ValidationResult(is_valid=False, error_message="Insufficient balance")
|
||||
|
||||
# Check risk limits
|
||||
current_risk = await self.risk_calculator.calculate_trade_risk(portfolio, trade_request)
|
||||
|
||||
if current_risk > portfolio.risk_tolerance:
|
||||
return ValidationResult(is_valid=False, error_message="Trade exceeds risk tolerance")
|
||||
|
||||
return ValidationResult(is_valid=True)
|
||||
|
||||
def _calculate_buy_amount(self, sell_amount: float, sell_price: float, buy_price: float) -> float:
|
||||
"""Calculate expected buy amount"""
|
||||
sell_value = sell_amount * sell_price
|
||||
return sell_value / buy_price
|
||||
|
||||
async def _update_portfolio_assets(self, portfolio: AgentPortfolio, trade: PortfolioTrade) -> None:
|
||||
"""Update portfolio assets after trade"""
|
||||
|
||||
# Update sell asset
|
||||
sell_asset = self.session.execute(
|
||||
select(PortfolioAsset).where(
|
||||
PortfolioAsset.portfolio_id == portfolio.id, PortfolioAsset.token_symbol == trade.sell_token
|
||||
)
|
||||
).first()
|
||||
|
||||
if sell_asset:
|
||||
sell_asset.balance -= trade.sell_amount
|
||||
sell_asset.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
# Update buy asset
|
||||
buy_asset = self.session.execute(
|
||||
select(PortfolioAsset).where(
|
||||
PortfolioAsset.portfolio_id == portfolio.id, PortfolioAsset.token_symbol == trade.buy_token
|
||||
)
|
||||
).first()
|
||||
|
||||
if buy_asset:
|
||||
buy_asset.balance += trade.buy_amount
|
||||
buy_asset.updated_at = datetime.now(timezone.utc)
|
||||
else:
|
||||
# Create new asset if it doesn't exist
|
||||
new_asset = PortfolioAsset(
|
||||
portfolio_id=portfolio.id,
|
||||
token_symbol=trade.buy_token,
|
||||
target_allocation=0.0,
|
||||
current_allocation=0.0,
|
||||
balance=trade.buy_amount,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
)
|
||||
self.session.add(new_asset)
|
||||
|
||||
async def _update_portfolio_metrics(self, portfolio: AgentPortfolio) -> None:
|
||||
"""Update portfolio value and allocations"""
|
||||
|
||||
portfolio_value = await self._calculate_portfolio_value(portfolio)
|
||||
|
||||
# Update current allocations
|
||||
assets = self.session.execute(select(PortfolioAsset).where(PortfolioAsset.portfolio_id == portfolio.id)).all()
|
||||
|
||||
for asset in assets:
|
||||
if asset.balance > 0:
|
||||
price = await self.price_service.get_price(asset.token_symbol)
|
||||
asset_value = asset.balance * price
|
||||
asset.current_allocation = (asset_value / portfolio_value) * 100
|
||||
asset.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
portfolio.total_value = portfolio_value
|
||||
portfolio.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
async def _calculate_portfolio_value(self, portfolio: AgentPortfolio) -> float:
|
||||
"""Calculate total portfolio value"""
|
||||
|
||||
assets = self.session.execute(select(PortfolioAsset).where(PortfolioAsset.portfolio_id == portfolio.id)).all()
|
||||
|
||||
total_value = 0.0
|
||||
for asset in assets:
|
||||
if asset.balance > 0:
|
||||
price = await self.price_service.get_price(asset.token_symbol)
|
||||
total_value += asset.balance * price
|
||||
|
||||
return total_value
|
||||
|
||||
async def _needs_rebalancing(self, portfolio: AgentPortfolio) -> bool:
|
||||
"""Check if portfolio needs rebalancing"""
|
||||
|
||||
# Check time-based rebalancing
|
||||
strategy = self.session.get(PortfolioStrategy, portfolio.strategy_id)
|
||||
if not strategy:
|
||||
return False
|
||||
|
||||
time_since_rebalance = datetime.now(timezone.utc) - portfolio.last_rebalance
|
||||
if time_since_rebalance > timedelta(seconds=strategy.rebalance_frequency):
|
||||
return True
|
||||
|
||||
# Check threshold-based rebalancing
|
||||
assets = self.session.execute(select(PortfolioAsset).where(PortfolioAsset.portfolio_id == portfolio.id)).all()
|
||||
|
||||
for asset in assets:
|
||||
if asset.balance > 0:
|
||||
deviation = abs(asset.current_allocation - asset.target_allocation)
|
||||
if deviation > 5.0: # 5% deviation threshold
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def _generate_rebalance_trades(
|
||||
self, portfolio: AgentPortfolio, optimal_allocations: dict[str, float]
|
||||
) -> list[TradeRequest]:
|
||||
"""Generate rebalancing trades"""
|
||||
|
||||
trades = []
|
||||
assets = self.session.execute(select(PortfolioAsset).where(PortfolioAsset.portfolio_id == portfolio.id)).all()
|
||||
|
||||
# Calculate current vs target allocations
|
||||
for asset in assets:
|
||||
target_allocation = optimal_allocations.get(asset.token_symbol, 0.0)
|
||||
current_allocation = asset.current_allocation
|
||||
|
||||
if abs(current_allocation - target_allocation) > 1.0: # 1% minimum deviation
|
||||
if current_allocation > target_allocation:
|
||||
# Sell excess
|
||||
excess_percentage = current_allocation - target_allocation
|
||||
sell_amount = (asset.balance * excess_percentage) / 100
|
||||
|
||||
# Find asset to buy
|
||||
for other_asset in assets:
|
||||
other_target = optimal_allocations.get(other_asset.token_symbol, 0.0)
|
||||
other_current = other_asset.current_allocation
|
||||
|
||||
if other_current < other_target:
|
||||
trade = TradeRequest(
|
||||
sell_token=asset.token_symbol,
|
||||
buy_token=other_asset.token_symbol,
|
||||
sell_amount=sell_amount,
|
||||
min_buy_amount=0, # Will be calculated during execution
|
||||
)
|
||||
trades.append(trade)
|
||||
break
|
||||
|
||||
return trades
|
||||
|
||||
async def _calculate_performance_metrics(self, portfolio: AgentPortfolio, period: str) -> dict:
|
||||
"""Calculate portfolio performance metrics"""
|
||||
|
||||
# Get historical trades
|
||||
trades = self.session.execute(
|
||||
select(PortfolioTrade)
|
||||
.where(PortfolioTrade.portfolio_id == portfolio.id)
|
||||
.order_by(PortfolioTrade.executed_at.desc())
|
||||
).all()
|
||||
|
||||
# Calculate returns, volatility, etc.
|
||||
# This is a simplified implementation
|
||||
current_value = await self._calculate_portfolio_value(portfolio)
|
||||
initial_value = portfolio.initial_capital
|
||||
|
||||
total_return = ((current_value - initial_value) / initial_value) * 100
|
||||
|
||||
return {
|
||||
"total_return": total_return,
|
||||
"current_value": current_value,
|
||||
"initial_value": initial_value,
|
||||
"total_trades": len(trades),
|
||||
"last_updated": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
|
||||
|
||||
class ValidationResult:
|
||||
"""Validation result for trade requests"""
|
||||
|
||||
def __init__(self, is_valid: bool, error_message: str = ""):
|
||||
self.is_valid = is_valid
|
||||
self.error_message = error_message
|
||||
903
apps/agent-management/src/app/services/agent_security.py
Executable file
903
apps/agent-management/src/app/services/agent_security.py
Executable file
@@ -0,0 +1,903 @@
|
||||
"""
|
||||
Agent Security and Audit Framework for Verifiable AI Agent Orchestration
|
||||
Implements comprehensive security, auditing, and trust establishment for agent executions
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
from datetime import datetime, timezone
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlmodel import JSON, Column, Field, Session, SQLModel, select
|
||||
|
||||
from app.domain.agent import AIAgentWorkflow, VerificationLevel
|
||||
|
||||
|
||||
class SecurityLevel(StrEnum):
|
||||
"""Security classification levels for agent operations"""
|
||||
|
||||
PUBLIC = "public"
|
||||
INTERNAL = "internal"
|
||||
CONFIDENTIAL = "confidential"
|
||||
RESTRICTED = "restricted"
|
||||
|
||||
|
||||
class AuditEventType(StrEnum):
|
||||
"""Types of audit events for agent operations"""
|
||||
|
||||
WORKFLOW_CREATED = "workflow_created"
|
||||
WORKFLOW_UPDATED = "workflow_updated"
|
||||
WORKFLOW_DELETED = "workflow_deleted"
|
||||
EXECUTION_STARTED = "execution_started"
|
||||
EXECUTION_COMPLETED = "execution_completed"
|
||||
EXECUTION_FAILED = "execution_failed"
|
||||
EXECUTION_CANCELLED = "execution_cancelled"
|
||||
STEP_STARTED = "step_started"
|
||||
STEP_COMPLETED = "step_completed"
|
||||
STEP_FAILED = "step_failed"
|
||||
VERIFICATION_COMPLETED = "verification_completed"
|
||||
VERIFICATION_FAILED = "verification_failed"
|
||||
SECURITY_VIOLATION = "security_violation"
|
||||
ACCESS_DENIED = "access_denied"
|
||||
SANDBOX_BREACH = "sandbox_breach"
|
||||
|
||||
|
||||
class AgentAuditLog(SQLModel, table=True):
|
||||
"""Comprehensive audit log for agent operations"""
|
||||
|
||||
__tablename__ = "agent_audit_logs"
|
||||
|
||||
id: str = Field(default_factory=lambda: f"audit_{uuid4().hex[:12]}", primary_key=True)
|
||||
|
||||
# Event information
|
||||
event_type: AuditEventType = Field(index=True)
|
||||
timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
|
||||
|
||||
# Entity references
|
||||
workflow_id: str | None = Field(index=True)
|
||||
execution_id: str | None = Field(index=True)
|
||||
step_id: str | None = Field(index=True)
|
||||
user_id: str | None = Field(index=True)
|
||||
|
||||
# Security context
|
||||
security_level: SecurityLevel = Field(default=SecurityLevel.PUBLIC)
|
||||
ip_address: str | None = Field(default=None)
|
||||
user_agent: str | None = Field(default=None)
|
||||
|
||||
# Event data
|
||||
event_data: dict[str, Any] = Field(default_factory=dict, sa_column=Column(JSON))
|
||||
previous_state: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON))
|
||||
new_state: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON))
|
||||
|
||||
# Security metadata
|
||||
risk_score: int = Field(default=0) # 0-100 risk assessment
|
||||
requires_investigation: bool = Field(default=False)
|
||||
investigation_notes: str | None = Field(default=None)
|
||||
|
||||
# Verification
|
||||
cryptographic_hash: str | None = Field(default=None)
|
||||
signature_valid: bool | None = Field(default=None)
|
||||
|
||||
# Metadata
|
||||
created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
class AgentSecurityPolicy(SQLModel, table=True):
|
||||
"""Security policies for agent operations"""
|
||||
|
||||
__tablename__ = "agent_security_policies"
|
||||
|
||||
id: str = Field(default_factory=lambda: f"policy_{uuid4().hex[:8]}", primary_key=True)
|
||||
|
||||
# Policy definition
|
||||
name: str = Field(max_length=100, unique=True)
|
||||
description: str = Field(default="")
|
||||
security_level: SecurityLevel = Field(default=SecurityLevel.PUBLIC)
|
||||
|
||||
# Policy rules
|
||||
allowed_step_types: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
max_execution_time: int = Field(default=3600) # seconds
|
||||
max_memory_usage: int = Field(default=8192) # MB
|
||||
require_verification: bool = Field(default=True)
|
||||
allowed_verification_levels: list[VerificationLevel] = Field(
|
||||
default_factory=lambda: [VerificationLevel.BASIC], sa_column=Column(JSON)
|
||||
)
|
||||
|
||||
# Resource limits
|
||||
max_concurrent_executions: int = Field(default=10)
|
||||
max_workflow_steps: int = Field(default=100)
|
||||
max_data_size: int = Field(default=1024 * 1024 * 1024) # 1GB
|
||||
|
||||
# Security requirements
|
||||
require_sandbox: bool = Field(default=False)
|
||||
require_audit_logging: bool = Field(default=True)
|
||||
require_encryption: bool = Field(default=False)
|
||||
|
||||
# Compliance
|
||||
compliance_standards: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
|
||||
# Status
|
||||
is_active: bool = Field(default=True)
|
||||
created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
class AgentTrustScore(SQLModel, table=True):
|
||||
"""Trust and reputation scoring for agents and users"""
|
||||
|
||||
__tablename__ = "agent_trust_scores"
|
||||
|
||||
id: str = Field(default_factory=lambda: f"trust_{uuid4().hex[:8]}", primary_key=True)
|
||||
|
||||
# Entity information
|
||||
entity_type: str = Field(index=True) # "agent", "user", "workflow"
|
||||
entity_id: str = Field(index=True)
|
||||
|
||||
# Trust metrics
|
||||
trust_score: float = Field(default=0.0, index=True) # 0-100
|
||||
reputation_score: float = Field(default=0.0) # 0-100
|
||||
|
||||
# Performance metrics
|
||||
total_executions: int = Field(default=0)
|
||||
successful_executions: int = Field(default=0)
|
||||
failed_executions: int = Field(default=0)
|
||||
verification_success_rate: float = Field(default=0.0)
|
||||
|
||||
# Security metrics
|
||||
security_violations: int = Field(default=0)
|
||||
policy_violations: int = Field(default=0)
|
||||
sandbox_breaches: int = Field(default=0)
|
||||
|
||||
# Time-based metrics
|
||||
last_execution: datetime | None = Field(default=None)
|
||||
last_violation: datetime | None = Field(default=None)
|
||||
average_execution_time: float | None = Field(default=None)
|
||||
|
||||
# Historical data
|
||||
execution_history: list[dict[str, Any]] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
violation_history: list[dict[str, Any]] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
|
||||
# Metadata
|
||||
created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
class AgentSandboxConfig(SQLModel, table=True):
|
||||
"""Sandboxing configuration for agent execution"""
|
||||
|
||||
__tablename__ = "agent_sandbox_configs"
|
||||
|
||||
id: str = Field(default_factory=lambda: f"sandbox_{uuid4().hex[:8]}", primary_key=True)
|
||||
|
||||
# Sandbox type
|
||||
sandbox_type: str = Field(default="process") # vm, process, none
|
||||
security_level: SecurityLevel = Field(default=SecurityLevel.PUBLIC)
|
||||
|
||||
# Resource limits
|
||||
cpu_limit: float = Field(default=1.0) # CPU cores
|
||||
memory_limit: int = Field(default=1024) # MB
|
||||
disk_limit: int = Field(default=10240) # MB
|
||||
network_access: bool = Field(default=False)
|
||||
|
||||
# Security restrictions
|
||||
allowed_commands: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
blocked_commands: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
allowed_file_paths: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
blocked_file_paths: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
|
||||
# Network restrictions
|
||||
allowed_domains: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
blocked_domains: list[str] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
allowed_ports: list[int] = Field(default_factory=list, sa_column=Column(JSON))
|
||||
|
||||
# Time limits
|
||||
max_execution_time: int = Field(default=3600) # seconds
|
||||
idle_timeout: int = Field(default=300) # seconds
|
||||
|
||||
# Monitoring
|
||||
enable_monitoring: bool = Field(default=True)
|
||||
log_all_commands: bool = Field(default=False)
|
||||
log_file_access: bool = Field(default=True)
|
||||
log_network_access: bool = Field(default=True)
|
||||
|
||||
# Status
|
||||
is_active: bool = Field(default=True)
|
||||
created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
class AgentAuditor:
|
||||
"""Comprehensive auditing system for agent operations"""
|
||||
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
self.security_policies = {}
|
||||
self.trust_manager = AgentTrustManager(session)
|
||||
self.sandbox_manager = AgentSandboxManager(session)
|
||||
|
||||
async def log_event(
|
||||
self,
|
||||
event_type: AuditEventType,
|
||||
workflow_id: str | None = None,
|
||||
execution_id: str | None = None,
|
||||
step_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
security_level: SecurityLevel = SecurityLevel.PUBLIC,
|
||||
event_data: dict[str, Any] | None = None,
|
||||
previous_state: dict[str, Any] | None = None,
|
||||
new_state: dict[str, Any] | None = None,
|
||||
ip_address: str | None = None,
|
||||
user_agent: str | None = None,
|
||||
) -> AgentAuditLog:
|
||||
"""Log an audit event with comprehensive security context"""
|
||||
|
||||
# Calculate risk score
|
||||
risk_score = self._calculate_risk_score(event_type, event_data, security_level)
|
||||
|
||||
# Create audit log entry
|
||||
audit_log = AgentAuditLog(
|
||||
event_type=event_type,
|
||||
workflow_id=workflow_id,
|
||||
execution_id=execution_id,
|
||||
step_id=step_id,
|
||||
user_id=user_id,
|
||||
security_level=security_level,
|
||||
ip_address=ip_address,
|
||||
user_agent=user_agent,
|
||||
event_data=event_data or {},
|
||||
previous_state=previous_state,
|
||||
new_state=new_state,
|
||||
risk_score=risk_score,
|
||||
requires_investigation=risk_score >= 70,
|
||||
cryptographic_hash=self._generate_event_hash(event_data),
|
||||
signature_valid=self._verify_signature(event_data),
|
||||
)
|
||||
|
||||
# Store audit log
|
||||
self.session.add(audit_log)
|
||||
self.session.commit()
|
||||
self.session.refresh(audit_log)
|
||||
|
||||
# Handle high-risk events
|
||||
if audit_log.requires_investigation:
|
||||
await self._handle_high_risk_event(audit_log)
|
||||
|
||||
logger.info(f"Audit event logged: {event_type.value} for workflow {workflow_id} execution {execution_id}")
|
||||
return audit_log
|
||||
|
||||
def _calculate_risk_score(
|
||||
self, event_type: AuditEventType, event_data: dict[str, Any], security_level: SecurityLevel
|
||||
) -> int:
|
||||
"""Calculate risk score for audit event"""
|
||||
|
||||
base_score = 0
|
||||
|
||||
# Event type risk
|
||||
event_risk_scores = {
|
||||
AuditEventType.SECURITY_VIOLATION: 90,
|
||||
AuditEventType.SANDBOX_BREACH: 85,
|
||||
AuditEventType.ACCESS_DENIED: 70,
|
||||
AuditEventType.VERIFICATION_FAILED: 50,
|
||||
AuditEventType.EXECUTION_FAILED: 30,
|
||||
AuditEventType.STEP_FAILED: 20,
|
||||
AuditEventType.EXECUTION_CANCELLED: 15,
|
||||
AuditEventType.WORKFLOW_DELETED: 10,
|
||||
AuditEventType.WORKFLOW_CREATED: 5,
|
||||
AuditEventType.EXECUTION_STARTED: 3,
|
||||
AuditEventType.EXECUTION_COMPLETED: 1,
|
||||
AuditEventType.STEP_STARTED: 1,
|
||||
AuditEventType.STEP_COMPLETED: 1,
|
||||
AuditEventType.VERIFICATION_COMPLETED: 1,
|
||||
}
|
||||
|
||||
base_score += event_risk_scores.get(event_type, 0)
|
||||
|
||||
# Security level adjustment
|
||||
security_multipliers = {
|
||||
SecurityLevel.PUBLIC: 1.0,
|
||||
SecurityLevel.INTERNAL: 1.2,
|
||||
SecurityLevel.CONFIDENTIAL: 1.5,
|
||||
SecurityLevel.RESTRICTED: 2.0,
|
||||
}
|
||||
|
||||
base_score = int(base_score * security_multipliers[security_level])
|
||||
|
||||
# Event data analysis
|
||||
if event_data:
|
||||
# Check for suspicious patterns
|
||||
if event_data.get("error_message"):
|
||||
base_score += 10
|
||||
if event_data.get("execution_time", 0) > 3600: # > 1 hour
|
||||
base_score += 5
|
||||
if event_data.get("memory_usage", 0) > 8192: # > 8GB
|
||||
base_score += 5
|
||||
|
||||
return min(base_score, 100)
|
||||
|
||||
def _generate_event_hash(self, event_data: dict[str, Any]) -> str:
|
||||
"""Generate cryptographic hash for event data"""
|
||||
if not event_data:
|
||||
return None
|
||||
|
||||
# Create canonical JSON representation
|
||||
canonical_json = json.dumps(event_data, sort_keys=True, separators=(",", ":"))
|
||||
return hashlib.sha256(canonical_json.encode()).hexdigest()
|
||||
|
||||
def _verify_signature(self, event_data: dict[str, Any]) -> bool | None:
|
||||
"""Verify cryptographic signature of event data
|
||||
|
||||
Note: Full signature verification requires:
|
||||
1. Extract signature from event_data
|
||||
2. Verify against expected public key
|
||||
3. Use appropriate crypto library (e.g., cryptography, eth_keys)
|
||||
Currently returns None (not verified) for compatibility.
|
||||
"""
|
||||
try:
|
||||
# Check if signature data exists
|
||||
if "signature" not in event_data or "public_key" not in event_data:
|
||||
return None
|
||||
|
||||
# Placeholder for actual signature verification
|
||||
# In production, use cryptography library to verify signature
|
||||
# from cryptography.hazmat.primitives import hashes
|
||||
# from cryptography.hazmat.primitives.asymmetric import padding
|
||||
|
||||
# For now, return None to indicate not verified
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Signature verification failed: {e}")
|
||||
return False
|
||||
|
||||
async def _handle_high_risk_event(self, audit_log: AgentAuditLog):
|
||||
"""Handle high-risk audit events requiring investigation"""
|
||||
|
||||
logger.warning(f"High-risk audit event detected: {audit_log.event_type.value} (Score: {audit_log.risk_score})")
|
||||
|
||||
# Create investigation record
|
||||
investigation_notes = f"High-risk event detected on {audit_log.timestamp}. "
|
||||
investigation_notes += f"Event type: {audit_log.event_type.value}, "
|
||||
investigation_notes += f"Risk score: {audit_log.risk_score}. "
|
||||
investigation_notes += "Requires manual investigation."
|
||||
|
||||
# Update audit log
|
||||
audit_log.investigation_notes = investigation_notes
|
||||
audit_log.investigation_status = "pending"
|
||||
audit_log.investigation_required = True
|
||||
self.session.commit()
|
||||
|
||||
# Send alert to security team (placeholder for actual alerting system)
|
||||
# In production, integrate with email, Slack, or other alerting systems
|
||||
logger.critical(f"SECURITY ALERT: High-risk event requires investigation - Event ID: {audit_log.id}")
|
||||
|
||||
# Create investigation ticket (placeholder for ticketing system integration)
|
||||
# In production, integrate with Jira, GitHub Issues, or other ticketing systems
|
||||
logger.info(f"Investigation ticket would be created for event: {audit_log.id}")
|
||||
|
||||
# Temporarily suspend related entities if needed (placeholder for suspension logic)
|
||||
# In production, implement suspension logic based on risk level and event type
|
||||
if audit_log.risk_score >= 0.9:
|
||||
logger.warning(f"Critical risk score ({audit_log.risk_score}) - entity suspension recommended")
|
||||
# Placeholder for actual suspension logic
|
||||
# await self._suspend_entity_if_needed(audit_log)
|
||||
|
||||
|
||||
class AgentTrustManager:
|
||||
"""Trust and reputation management for agents and users"""
|
||||
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
|
||||
async def update_trust_score(
|
||||
self,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
execution_success: bool,
|
||||
execution_time: float | None = None,
|
||||
security_violation: bool = False,
|
||||
policy_violation: bool = bool,
|
||||
) -> AgentTrustScore:
|
||||
"""Update trust score based on execution results"""
|
||||
|
||||
# Get or create trust score record
|
||||
trust_score = self.session.execute(
|
||||
select(AgentTrustScore).where(
|
||||
(AgentTrustScore.entity_type == entity_type) & (AgentTrustScore.entity_id == entity_id)
|
||||
)
|
||||
).first()
|
||||
|
||||
if not trust_score:
|
||||
trust_score = AgentTrustScore(entity_type=entity_type, entity_id=entity_id)
|
||||
self.session.add(trust_score)
|
||||
|
||||
# Update metrics
|
||||
trust_score.total_executions += 1
|
||||
|
||||
if execution_success:
|
||||
trust_score.successful_executions += 1
|
||||
else:
|
||||
trust_score.failed_executions += 1
|
||||
|
||||
if security_violation:
|
||||
trust_score.security_violations += 1
|
||||
trust_score.last_violation = datetime.now(timezone.utc)
|
||||
trust_score.violation_history.append({"timestamp": datetime.now(timezone.utc).isoformat(), "type": "security_violation"})
|
||||
|
||||
if policy_violation:
|
||||
trust_score.policy_violations += 1
|
||||
trust_score.last_violation = datetime.now(timezone.utc)
|
||||
trust_score.violation_history.append({"timestamp": datetime.now(timezone.utc).isoformat(), "type": "policy_violation"})
|
||||
|
||||
# Calculate scores
|
||||
trust_score.trust_score = self._calculate_trust_score(trust_score)
|
||||
trust_score.reputation_score = self._calculate_reputation_score(trust_score)
|
||||
trust_score.verification_success_rate = (
|
||||
trust_score.successful_executions / trust_score.total_executions * 100 if trust_score.total_executions > 0 else 0
|
||||
)
|
||||
|
||||
# Update execution metrics
|
||||
if execution_time:
|
||||
if trust_score.average_execution_time is None:
|
||||
trust_score.average_execution_time = execution_time
|
||||
else:
|
||||
trust_score.average_execution_time = (
|
||||
trust_score.average_execution_time * (trust_score.total_executions - 1) + execution_time
|
||||
) / trust_score.total_executions
|
||||
|
||||
trust_score.last_execution = datetime.now(timezone.utc)
|
||||
trust_score.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
self.session.commit()
|
||||
self.session.refresh(trust_score)
|
||||
|
||||
return trust_score
|
||||
|
||||
def _calculate_trust_score(self, trust_score: AgentTrustScore) -> float:
|
||||
"""Calculate overall trust score"""
|
||||
|
||||
base_score = 50.0 # Start at neutral
|
||||
|
||||
# Success rate impact
|
||||
if trust_score.total_executions > 0:
|
||||
success_rate = trust_score.successful_executions / trust_score.total_executions
|
||||
base_score += (success_rate - 0.5) * 40 # +/- 20 points
|
||||
|
||||
# Security violations penalty
|
||||
violation_penalty = trust_score.security_violations * 10
|
||||
base_score -= violation_penalty
|
||||
|
||||
# Policy violations penalty
|
||||
policy_penalty = trust_score.policy_violations * 5
|
||||
base_score -= policy_penalty
|
||||
|
||||
# Recency bonus (recent successful executions)
|
||||
if trust_score.last_execution:
|
||||
days_since_last = (datetime.now(timezone.utc) - trust_score.last_execution).days
|
||||
if days_since_last < 7:
|
||||
base_score += 5 # Recent activity bonus
|
||||
elif days_since_last > 30:
|
||||
base_score -= 10 # Inactivity penalty
|
||||
|
||||
return max(0.0, min(100.0, base_score))
|
||||
|
||||
def _calculate_reputation_score(self, trust_score: AgentTrustScore) -> float:
|
||||
"""Calculate reputation score based on long-term performance"""
|
||||
|
||||
base_score = 50.0
|
||||
|
||||
# Long-term success rate
|
||||
if trust_score.total_executions >= 10:
|
||||
success_rate = trust_score.successful_executions / trust_score.total_executions
|
||||
base_score += (success_rate - 0.5) * 30 # +/- 15 points
|
||||
|
||||
# Volume bonus (more executions = more data points)
|
||||
volume_bonus = min(trust_score.total_executions / 100, 10) # Max 10 points
|
||||
base_score += volume_bonus
|
||||
|
||||
# Security record
|
||||
if trust_score.security_violations == 0 and trust_score.policy_violations == 0:
|
||||
base_score += 10 # Clean record bonus
|
||||
else:
|
||||
violation_penalty = (trust_score.security_violations + trust_score.policy_violations) * 2
|
||||
base_score -= violation_penalty
|
||||
|
||||
return max(0.0, min(100.0, base_score))
|
||||
|
||||
|
||||
class AgentSandboxManager:
|
||||
"""Sandboxing and isolation management for agent execution"""
|
||||
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
|
||||
async def create_sandbox_environment(
|
||||
self,
|
||||
execution_id: str,
|
||||
security_level: SecurityLevel = SecurityLevel.PUBLIC,
|
||||
workflow_requirements: dict[str, Any] | None = None,
|
||||
) -> AgentSandboxConfig:
|
||||
"""Create sandbox environment for agent execution"""
|
||||
|
||||
# Get appropriate sandbox configuration
|
||||
sandbox_config = self._get_sandbox_config(security_level)
|
||||
|
||||
# Customize based on workflow requirements
|
||||
if workflow_requirements:
|
||||
sandbox_config = self._customize_sandbox(sandbox_config, workflow_requirements)
|
||||
|
||||
# Create sandbox record
|
||||
sandbox = AgentSandboxConfig(
|
||||
id=f"sandbox_{execution_id}",
|
||||
sandbox_type=sandbox_config["type"],
|
||||
security_level=security_level,
|
||||
cpu_limit=sandbox_config["cpu_limit"],
|
||||
memory_limit=sandbox_config["memory_limit"],
|
||||
disk_limit=sandbox_config["disk_limit"],
|
||||
network_access=sandbox_config["network_access"],
|
||||
allowed_commands=sandbox_config["allowed_commands"],
|
||||
blocked_commands=sandbox_config["blocked_commands"],
|
||||
allowed_file_paths=sandbox_config["allowed_file_paths"],
|
||||
blocked_file_paths=sandbox_config["blocked_file_paths"],
|
||||
allowed_domains=sandbox_config["allowed_domains"],
|
||||
blocked_domains=sandbox_config["blocked_domains"],
|
||||
allowed_ports=sandbox_config["allowed_ports"],
|
||||
max_execution_time=sandbox_config["max_execution_time"],
|
||||
idle_timeout=sandbox_config["idle_timeout"],
|
||||
enable_monitoring=sandbox_config["enable_monitoring"],
|
||||
log_all_commands=sandbox_config["log_all_commands"],
|
||||
log_file_access=sandbox_config["log_file_access"],
|
||||
log_network_access=sandbox_config["log_network_access"],
|
||||
)
|
||||
|
||||
self.session.add(sandbox)
|
||||
self.session.commit()
|
||||
self.session.refresh(sandbox)
|
||||
|
||||
# Sandbox environment creation requires integration with:
|
||||
# 1. Podman for container isolation
|
||||
# 2. Firecracker/gVisor for VM-level isolation
|
||||
# 3. Process isolation using seccomp, namespaces
|
||||
# 4. Network isolation using virtual networks
|
||||
# Currently storing configuration only - actual sandbox creation
|
||||
# would be implemented by the execution orchestrator.
|
||||
|
||||
logger.info(f"Created sandbox configuration for execution {execution_id}")
|
||||
return sandbox
|
||||
|
||||
def _get_sandbox_config(self, security_level: SecurityLevel) -> dict[str, Any]:
|
||||
"""Get sandbox configuration based on security level"""
|
||||
|
||||
configs = {
|
||||
SecurityLevel.PUBLIC: {
|
||||
"type": "process",
|
||||
"cpu_limit": 1.0,
|
||||
"memory_limit": 1024,
|
||||
"disk_limit": 10240,
|
||||
"network_access": False,
|
||||
"allowed_commands": ["python", "node", "java"],
|
||||
"blocked_commands": ["rm", "sudo", "chmod", "chown"],
|
||||
"allowed_file_paths": ["/tmp", "/workspace"],
|
||||
"blocked_file_paths": ["/etc", "/root", "/home"],
|
||||
"allowed_domains": [],
|
||||
"blocked_domains": [],
|
||||
"allowed_ports": [],
|
||||
"max_execution_time": 3600,
|
||||
"idle_timeout": 300,
|
||||
"enable_monitoring": True,
|
||||
"log_all_commands": False,
|
||||
"log_file_access": True,
|
||||
"log_network_access": True,
|
||||
},
|
||||
SecurityLevel.INTERNAL: {
|
||||
"type": "docker",
|
||||
"cpu_limit": 2.0,
|
||||
"memory_limit": 2048,
|
||||
"disk_limit": 20480,
|
||||
"network_access": True,
|
||||
"allowed_commands": ["python", "node", "java", "curl", "wget"],
|
||||
"blocked_commands": ["rm", "sudo", "chmod", "chown", "iptables"],
|
||||
"allowed_file_paths": ["/tmp", "/workspace", "/app"],
|
||||
"blocked_file_paths": ["/etc", "/root", "/home", "/var"],
|
||||
"allowed_domains": ["*.internal.com", "*.api.internal"],
|
||||
"blocked_domains": ["malicious.com", "*.suspicious.net"],
|
||||
"allowed_ports": [80, 443, 8000, 8001, 8002, 8003, 8010, 8011, 8012, 8013, 8014, 8015, 8016],
|
||||
"max_execution_time": 7200,
|
||||
"idle_timeout": 600,
|
||||
"enable_monitoring": True,
|
||||
"log_all_commands": True,
|
||||
"log_file_access": True,
|
||||
"log_network_access": True,
|
||||
},
|
||||
SecurityLevel.CONFIDENTIAL: {
|
||||
"type": "docker",
|
||||
"cpu_limit": 4.0,
|
||||
"memory_limit": 4096,
|
||||
"disk_limit": 40960,
|
||||
"network_access": True,
|
||||
"allowed_commands": ["python", "node", "java", "curl", "wget", "git"],
|
||||
"blocked_commands": ["rm", "sudo", "chmod", "chown", "iptables", "systemctl"],
|
||||
"allowed_file_paths": ["/tmp", "/workspace", "/app", "/data"],
|
||||
"blocked_file_paths": ["/etc", "/root", "/home", "/var", "/sys", "/proc"],
|
||||
"allowed_domains": ["*.internal.com", "*.api.internal", "*.trusted.com"],
|
||||
"blocked_domains": ["malicious.com", "*.suspicious.net", "*.evil.org"],
|
||||
"allowed_ports": [80, 443, 8000, 8001, 8002, 8003, 8010, 8011, 8012, 8013, 8014, 8015, 8016],
|
||||
"max_execution_time": 14400,
|
||||
"idle_timeout": 1800,
|
||||
"enable_monitoring": True,
|
||||
"log_all_commands": True,
|
||||
"log_file_access": True,
|
||||
"log_network_access": True,
|
||||
},
|
||||
SecurityLevel.RESTRICTED: {
|
||||
"type": "vm",
|
||||
"cpu_limit": 8.0,
|
||||
"memory_limit": 8192,
|
||||
"disk_limit": 81920,
|
||||
"network_access": True,
|
||||
"allowed_commands": ["python", "node", "java", "curl", "wget", "git", "docker"],
|
||||
"blocked_commands": ["rm", "sudo", "chmod", "chown", "iptables", "systemctl", "systemd"],
|
||||
"allowed_file_paths": ["/tmp", "/workspace", "/app", "/data", "/shared"],
|
||||
"blocked_file_paths": ["/etc", "/root", "/home", "/var", "/sys", "/proc", "/boot"],
|
||||
"allowed_domains": ["*.internal.com", "*.api.internal", "*.trusted.com", "*.partner.com"],
|
||||
"blocked_domains": ["malicious.com", "*.suspicious.net", "*.evil.org"],
|
||||
"allowed_ports": [80, 443, 8000, 8001, 8002, 8003, 8010, 8011, 8012, 8013, 8014, 8015, 8016, 22, 25],
|
||||
"max_execution_time": 28800,
|
||||
"idle_timeout": 3600,
|
||||
"enable_monitoring": True,
|
||||
"log_all_commands": True,
|
||||
"log_file_access": True,
|
||||
"log_network_access": True,
|
||||
},
|
||||
}
|
||||
|
||||
return configs.get(security_level, configs[SecurityLevel.PUBLIC])
|
||||
|
||||
def _customize_sandbox(self, base_config: dict[str, Any], requirements: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Customize sandbox configuration based on workflow requirements"""
|
||||
|
||||
config = base_config.copy()
|
||||
|
||||
# Adjust resources based on requirements
|
||||
if "cpu_cores" in requirements:
|
||||
config["cpu_limit"] = max(config["cpu_limit"], requirements["cpu_cores"])
|
||||
|
||||
if "memory_mb" in requirements:
|
||||
config["memory_limit"] = max(config["memory_limit"], requirements["memory_mb"])
|
||||
|
||||
if "disk_mb" in requirements:
|
||||
config["disk_limit"] = max(config["disk_limit"], requirements["disk_mb"])
|
||||
|
||||
if "max_execution_time" in requirements:
|
||||
config["max_execution_time"] = min(config["max_execution_time"], requirements["max_execution_time"])
|
||||
|
||||
# Add custom commands if specified
|
||||
if "allowed_commands" in requirements:
|
||||
config["allowed_commands"].extend(requirements["allowed_commands"])
|
||||
|
||||
if "blocked_commands" in requirements:
|
||||
config["blocked_commands"].extend(requirements["blocked_commands"])
|
||||
|
||||
# Add network access if required
|
||||
if "network_access" in requirements:
|
||||
config["network_access"] = config["network_access"] or requirements["network_access"]
|
||||
|
||||
return config
|
||||
|
||||
async def monitor_sandbox(self, execution_id: str) -> dict[str, Any]:
|
||||
"""Monitor sandbox execution for security violations
|
||||
|
||||
Note: Actual sandbox monitoring requires integration with:
|
||||
1. Container runtime metrics (Docker stats, containerd)
|
||||
2. Process monitoring (psutil, /proc filesystem)
|
||||
3. Network monitoring (iptables, eBPF)
|
||||
4. File system monitoring (inotify, auditd)
|
||||
Currently returning placeholder monitoring data.
|
||||
"""
|
||||
# Get sandbox configuration
|
||||
sandbox = self.session.execute(
|
||||
select(AgentSandboxConfig).where(AgentSandboxConfig.id == f"sandbox_{execution_id}")
|
||||
).first()
|
||||
|
||||
if not sandbox:
|
||||
raise ValueError(f"Sandbox not found for execution {execution_id}")
|
||||
|
||||
# Placeholder for actual monitoring implementation
|
||||
# In production, integrate with container runtime for real metrics
|
||||
monitoring_data = {
|
||||
"execution_id": execution_id,
|
||||
"sandbox_type": sandbox.sandbox_type,
|
||||
"security_level": sandbox.security_level,
|
||||
"resource_usage": {"cpu_percent": 0.0, "memory_mb": 0, "disk_mb": 0},
|
||||
"security_events": [],
|
||||
"command_count": 0,
|
||||
"file_access_count": 0,
|
||||
"network_access_count": 0,
|
||||
"status": "configured",
|
||||
"note": "Monitoring requires sandbox runtime integration"
|
||||
}
|
||||
|
||||
return monitoring_data
|
||||
|
||||
async def cleanup_sandbox(self, execution_id: str) -> bool:
|
||||
"""Clean up sandbox environment after execution"""
|
||||
|
||||
try:
|
||||
# Get sandbox record
|
||||
sandbox = self.session.execute(
|
||||
select(AgentSandboxConfig).where(AgentSandboxConfig.id == f"sandbox_{execution_id}")
|
||||
).first()
|
||||
|
||||
if sandbox:
|
||||
# Mark as inactive
|
||||
sandbox.is_active = False
|
||||
sandbox.updated_at = datetime.now(timezone.utc)
|
||||
self.session.commit()
|
||||
|
||||
# Sandbox cleanup requires integration with:
|
||||
# 1. Docker/Podman: docker stop/rm, podman stop/rm
|
||||
# 2. VM management: Firecracker terminate
|
||||
# 3. Process cleanup: kill processes, cleanup namespaces
|
||||
# 4. Resource cleanup: remove temp files, network interfaces
|
||||
# Currently marking as inactive - actual cleanup would be
|
||||
# implemented by the execution orchestrator.
|
||||
# Future implementation: await self._cleanup_docker_sandbox(sandbox)
|
||||
|
||||
logger.info(f"Marked sandbox as inactive for execution {execution_id}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to cleanup sandbox for execution {execution_id}: {e}")
|
||||
return False
|
||||
|
||||
|
||||
class AgentSecurityManager:
|
||||
"""Main security management interface for agent operations"""
|
||||
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
self.auditor = AgentAuditor(session)
|
||||
self.trust_manager = AgentTrustManager(session)
|
||||
self.sandbox_manager = AgentSandboxManager(session)
|
||||
|
||||
async def create_security_policy(
|
||||
self, name: str, description: str, security_level: SecurityLevel, policy_rules: dict[str, Any]
|
||||
) -> AgentSecurityPolicy:
|
||||
"""Create a new security policy"""
|
||||
|
||||
policy = AgentSecurityPolicy(name=name, description=description, security_level=security_level, **policy_rules)
|
||||
|
||||
self.session.add(policy)
|
||||
self.session.commit()
|
||||
self.session.refresh(policy)
|
||||
|
||||
# Log policy creation
|
||||
await self.auditor.log_event(
|
||||
AuditEventType.WORKFLOW_CREATED,
|
||||
user_id="system",
|
||||
security_level=SecurityLevel.INTERNAL,
|
||||
event_data={"policy_name": name, "policy_id": policy.id},
|
||||
new_state={"policy": policy.dict()},
|
||||
)
|
||||
|
||||
return policy
|
||||
|
||||
async def validate_workflow_security(self, workflow: AIAgentWorkflow, user_id: str) -> dict[str, Any]:
|
||||
"""Validate workflow against security policies"""
|
||||
|
||||
validation_result = {
|
||||
"valid": True,
|
||||
"violations": [],
|
||||
"warnings": [],
|
||||
"required_security_level": SecurityLevel.PUBLIC,
|
||||
"recommendations": [],
|
||||
}
|
||||
|
||||
# Check for security-sensitive operations
|
||||
security_sensitive_steps = []
|
||||
for step_data in workflow.steps.values():
|
||||
if step_data.get("step_type") in ["training", "data_processing"]:
|
||||
security_sensitive_steps.append(step_data.get("name"))
|
||||
|
||||
if security_sensitive_steps:
|
||||
validation_result["warnings"].append(f"Security-sensitive steps detected: {security_sensitive_steps}")
|
||||
validation_result["recommendations"].append(
|
||||
"Consider using higher security level for workflows with sensitive operations"
|
||||
)
|
||||
|
||||
# Check execution time
|
||||
if workflow.max_execution_time > 3600: # > 1 hour
|
||||
validation_result["warnings"].append(
|
||||
f"Long execution time ({workflow.max_execution_time}s) may require additional security measures"
|
||||
)
|
||||
|
||||
# Check verification requirements
|
||||
if not workflow.requires_verification:
|
||||
validation_result["violations"].append(
|
||||
"Workflow does not require verification - this is not recommended for production use"
|
||||
)
|
||||
validation_result["valid"] = False
|
||||
|
||||
# Determine required security level
|
||||
if workflow.requires_verification and workflow.verification_level == VerificationLevel.ZERO_KNOWLEDGE:
|
||||
validation_result["required_security_level"] = SecurityLevel.RESTRICTED
|
||||
elif workflow.requires_verification and workflow.verification_level == VerificationLevel.FULL:
|
||||
validation_result["required_security_level"] = SecurityLevel.CONFIDENTIAL
|
||||
elif workflow.requires_verification:
|
||||
validation_result["required_security_level"] = SecurityLevel.INTERNAL
|
||||
|
||||
# Log security validation
|
||||
await self.auditor.log_event(
|
||||
AuditEventType.WORKFLOW_CREATED,
|
||||
workflow_id=workflow.id,
|
||||
user_id=user_id,
|
||||
security_level=validation_result["required_security_level"],
|
||||
event_data={"validation_result": validation_result},
|
||||
)
|
||||
|
||||
return validation_result
|
||||
|
||||
async def monitor_execution_security(self, execution_id: str, workflow_id: str) -> dict[str, Any]:
|
||||
"""Monitor execution for security violations"""
|
||||
|
||||
monitoring_result = {
|
||||
"execution_id": execution_id,
|
||||
"workflow_id": workflow_id,
|
||||
"security_status": "monitoring",
|
||||
"violations": [],
|
||||
"alerts": [],
|
||||
}
|
||||
|
||||
try:
|
||||
# Monitor sandbox
|
||||
sandbox_monitoring = await self.sandbox_manager.monitor_sandbox(execution_id)
|
||||
|
||||
# Check for resource violations
|
||||
if sandbox_monitoring["resource_usage"]["cpu_percent"] > 90:
|
||||
monitoring_result["violations"].append("High CPU usage detected")
|
||||
monitoring_result["alerts"].append("CPU usage exceeded 90%")
|
||||
|
||||
if sandbox_monitoring["resource_usage"]["memory_mb"] > sandbox_monitoring["resource_usage"]["memory_mb"] * 0.9:
|
||||
monitoring_result["violations"].append("High memory usage detected")
|
||||
monitoring_result["alerts"].append("Memory usage exceeded 90% of limit")
|
||||
|
||||
# Check for security events
|
||||
if sandbox_monitoring["security_events"]:
|
||||
monitoring_result["violations"].extend(sandbox_monitoring["security_events"])
|
||||
monitoring_result["alerts"].extend(
|
||||
f"Security event: {event}" for event in sandbox_monitoring["security_events"]
|
||||
)
|
||||
|
||||
# Update security status
|
||||
if monitoring_result["violations"]:
|
||||
monitoring_result["security_status"] = "violations_detected"
|
||||
await self.auditor.log_event(
|
||||
AuditEventType.SECURITY_VIOLATION,
|
||||
execution_id=execution_id,
|
||||
workflow_id=workflow_id,
|
||||
security_level=SecurityLevel.INTERNAL,
|
||||
event_data={"violations": monitoring_result["violations"]},
|
||||
requires_investigation=len(monitoring_result["violations"]) > 0,
|
||||
)
|
||||
else:
|
||||
monitoring_result["security_status"] = "secure"
|
||||
|
||||
except Exception as e:
|
||||
monitoring_result["security_status"] = "monitoring_failed"
|
||||
monitoring_result["alerts"].append(f"Security monitoring failed: {e}")
|
||||
await self.auditor.log_event(
|
||||
AuditEventType.SECURITY_VIOLATION,
|
||||
execution_id=execution_id,
|
||||
workflow_id=workflow_id,
|
||||
security_level=SecurityLevel.INTERNAL,
|
||||
event_data={"error": str(e)},
|
||||
requires_investigation=True,
|
||||
)
|
||||
|
||||
return monitoring_result
|
||||
533
apps/agent-management/src/app/services/agent_service.py
Executable file
533
apps/agent-management/src/app/services/agent_service.py
Executable file
@@ -0,0 +1,533 @@
|
||||
"""
|
||||
AI Agent Service for Verifiable AI Agent Orchestration
|
||||
Implements core orchestration logic and state management for AI agent workflows
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from typing import Any
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from sqlmodel import Session, select, update
|
||||
|
||||
from app.domain.agent import (
|
||||
AgentExecution,
|
||||
AgentExecutionRequest,
|
||||
AgentExecutionResponse,
|
||||
AgentExecutionStatus,
|
||||
AgentStatus,
|
||||
AgentStep,
|
||||
AgentStepExecution,
|
||||
AIAgentWorkflow,
|
||||
StepType,
|
||||
VerificationLevel,
|
||||
)
|
||||
|
||||
|
||||
# Mock CoordinatorClient for now
|
||||
class CoordinatorClient:
|
||||
"""Mock coordinator client for agent orchestration"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AgentStateManager:
|
||||
"""Manages persistent state for AI agent executions"""
|
||||
|
||||
def __init__(self, session: Session):
|
||||
self.session = session
|
||||
|
||||
async def create_execution(
|
||||
self, workflow_id: str, client_id: str, verification_level: VerificationLevel = VerificationLevel.BASIC
|
||||
) -> AgentExecution:
|
||||
"""Create a new agent execution record"""
|
||||
|
||||
execution = AgentExecution(workflow_id=workflow_id, client_id=client_id, verification_level=verification_level)
|
||||
|
||||
self.session.add(execution)
|
||||
self.session.commit()
|
||||
self.session.refresh(execution)
|
||||
|
||||
logger.info(f"Created agent execution: {execution.id}")
|
||||
return execution
|
||||
|
||||
async def update_execution_status(self, execution_id: str, status: AgentStatus, **kwargs) -> AgentExecution:
|
||||
"""Update execution status and related fields"""
|
||||
|
||||
stmt = (
|
||||
update(AgentExecution)
|
||||
.where(AgentExecution.id == execution_id)
|
||||
.values(status=status, updated_at=datetime.now(timezone.utc), **kwargs)
|
||||
)
|
||||
|
||||
self.session.execute(stmt)
|
||||
self.session.commit()
|
||||
|
||||
# Get updated execution
|
||||
execution = self.session.get(AgentExecution, execution_id)
|
||||
logger.info(f"Updated execution {execution_id} status to {status}")
|
||||
return execution
|
||||
|
||||
async def get_execution(self, execution_id: str) -> AgentExecution | None:
|
||||
"""Get execution by ID"""
|
||||
return self.session.get(AgentExecution, execution_id)
|
||||
|
||||
async def get_workflow(self, workflow_id: str) -> AIAgentWorkflow | None:
|
||||
"""Get workflow by ID"""
|
||||
return self.session.get(AIAgentWorkflow, workflow_id)
|
||||
|
||||
async def get_workflow_steps(self, workflow_id: str) -> list[AgentStep]:
|
||||
"""Get all steps for a workflow"""
|
||||
stmt = select(AgentStep).where(AgentStep.workflow_id == workflow_id).order_by(AgentStep.step_order)
|
||||
return self.session.execute(stmt).all()
|
||||
|
||||
async def create_step_execution(self, execution_id: str, step_id: str) -> AgentStepExecution:
|
||||
"""Create a step execution record"""
|
||||
|
||||
step_execution = AgentStepExecution(execution_id=execution_id, step_id=step_id)
|
||||
|
||||
self.session.add(step_execution)
|
||||
self.session.commit()
|
||||
self.session.refresh(step_execution)
|
||||
|
||||
return step_execution
|
||||
|
||||
async def update_step_execution(self, step_execution_id: str, **kwargs) -> AgentStepExecution:
|
||||
"""Update step execution"""
|
||||
|
||||
stmt = (
|
||||
update(AgentStepExecution)
|
||||
.where(AgentStepExecution.id == step_execution_id)
|
||||
.values(updated_at=datetime.now(timezone.utc), **kwargs)
|
||||
)
|
||||
|
||||
self.session.execute(stmt)
|
||||
self.session.commit()
|
||||
|
||||
step_execution = self.session.get(AgentStepExecution, step_execution_id)
|
||||
return step_execution
|
||||
|
||||
|
||||
class AgentVerifier:
|
||||
"""Handles verification of agent executions"""
|
||||
|
||||
def __init__(self, cuda_accelerator=None):
|
||||
self.cuda_accelerator = cuda_accelerator
|
||||
|
||||
async def verify_step_execution(
|
||||
self, step_execution: AgentStepExecution, verification_level: VerificationLevel
|
||||
) -> dict[str, Any]:
|
||||
"""Verify a single step execution"""
|
||||
|
||||
verification_result = {
|
||||
"verified": False,
|
||||
"proof": None,
|
||||
"verification_time": 0.0,
|
||||
"verification_level": verification_level,
|
||||
}
|
||||
|
||||
try:
|
||||
if verification_level == VerificationLevel.ZERO_KNOWLEDGE:
|
||||
# Use ZK proof verification
|
||||
verification_result = await self._zk_verify_step(step_execution)
|
||||
elif verification_level == VerificationLevel.FULL:
|
||||
# Use comprehensive verification
|
||||
verification_result = await self._full_verify_step(step_execution)
|
||||
else:
|
||||
# Basic verification
|
||||
verification_result = await self._basic_verify_step(step_execution)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Step verification failed: {e}")
|
||||
verification_result["error"] = str(e)
|
||||
|
||||
return verification_result
|
||||
|
||||
async def _basic_verify_step(self, step_execution: AgentStepExecution) -> dict[str, Any]:
|
||||
"""Basic verification of step execution"""
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
# Basic checks: execution completed, has output, no errors
|
||||
verified = (
|
||||
step_execution.status == AgentStatus.COMPLETED
|
||||
and step_execution.output_data is not None
|
||||
and step_execution.error_message is None
|
||||
)
|
||||
|
||||
verification_time = (datetime.now(timezone.utc) - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"verified": verified,
|
||||
"proof": None,
|
||||
"verification_time": verification_time,
|
||||
"verification_level": VerificationLevel.BASIC,
|
||||
"checks": ["completion", "output_presence", "error_free"],
|
||||
}
|
||||
|
||||
async def _full_verify_step(self, step_execution: AgentStepExecution) -> dict[str, Any]:
|
||||
"""Full verification with additional checks"""
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
# Basic verification first
|
||||
basic_result = await self._basic_verify_step(step_execution)
|
||||
|
||||
if not basic_result["verified"]:
|
||||
return basic_result
|
||||
|
||||
# Additional checks: performance, resource usage
|
||||
additional_checks = []
|
||||
|
||||
# Check execution time is reasonable
|
||||
if step_execution.execution_time and step_execution.execution_time < 3600: # < 1 hour
|
||||
additional_checks.append("reasonable_execution_time")
|
||||
else:
|
||||
basic_result["verified"] = False
|
||||
|
||||
# Check memory usage
|
||||
if step_execution.memory_usage and step_execution.memory_usage < 8192: # < 8GB
|
||||
additional_checks.append("reasonable_memory_usage")
|
||||
|
||||
verification_time = (datetime.now(timezone.utc) - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"verified": basic_result["verified"],
|
||||
"proof": None,
|
||||
"verification_time": verification_time,
|
||||
"verification_level": VerificationLevel.FULL,
|
||||
"checks": basic_result["checks"] + additional_checks,
|
||||
}
|
||||
|
||||
async def _zk_verify_step(self, step_execution: AgentStepExecution) -> dict[str, Any]:
|
||||
"""Zero-knowledge proof verification
|
||||
|
||||
Note: Full ZK proof implementation requires integration with ZK-SNARKs/ZK-STARKs libraries.
|
||||
Currently using full verification as fallback. Future implementation should:
|
||||
1. Generate ZK proof from step execution
|
||||
2. Verify proof against public parameters
|
||||
3. Return verification result with proof hash
|
||||
"""
|
||||
datetime.now(timezone.utc)
|
||||
|
||||
# For now, fall back to full verification
|
||||
# ZK proof generation and verification requires specialized cryptographic libraries
|
||||
result = await self._full_verify_step(step_execution)
|
||||
result["verification_level"] = VerificationLevel.ZERO_KNOWLEDGE
|
||||
result["note"] = "ZK verification using full verification fallback (requires ZK-SNARKs integration)"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class AIAgentOrchestrator:
|
||||
"""Orchestrates execution of AI agent workflows"""
|
||||
|
||||
def __init__(self, session: Session, coordinator_client: CoordinatorClient):
|
||||
self.session = session
|
||||
self.coordinator = coordinator_client
|
||||
self.state_manager = AgentStateManager(session)
|
||||
self.verifier = AgentVerifier()
|
||||
|
||||
async def execute_workflow(self, request: AgentExecutionRequest, client_id: str) -> AgentExecutionResponse:
|
||||
"""Execute an AI agent workflow with verification"""
|
||||
|
||||
# Get workflow
|
||||
workflow = await self.state_manager.get_workflow(request.workflow_id)
|
||||
if not workflow:
|
||||
raise ValueError(f"Workflow not found: {request.workflow_id}")
|
||||
|
||||
# Create execution
|
||||
execution = await self.state_manager.create_execution(
|
||||
workflow_id=request.workflow_id, client_id=client_id, verification_level=request.verification_level
|
||||
)
|
||||
|
||||
try:
|
||||
# Start execution
|
||||
await self.state_manager.update_execution_status(
|
||||
execution.id, status=AgentStatus.RUNNING, started_at=datetime.now(timezone.utc), total_steps=len(workflow.steps)
|
||||
)
|
||||
|
||||
# Execute steps asynchronously
|
||||
asyncio.create_task(self._execute_steps_async(execution.id, request.inputs))
|
||||
|
||||
# Return initial response
|
||||
return AgentExecutionResponse(
|
||||
execution_id=execution.id,
|
||||
workflow_id=workflow.id,
|
||||
status=execution.status,
|
||||
current_step=0,
|
||||
total_steps=len(workflow.steps),
|
||||
started_at=execution.started_at,
|
||||
estimated_completion=self._estimate_completion(execution),
|
||||
current_cost=0.0,
|
||||
estimated_total_cost=self._estimate_cost(workflow),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
await self._handle_execution_failure(execution.id, e)
|
||||
raise
|
||||
|
||||
async def get_execution_status(self, execution_id: str) -> AgentExecutionStatus:
|
||||
"""Get current execution status"""
|
||||
|
||||
execution = await self.state_manager.get_execution(execution_id)
|
||||
if not execution:
|
||||
raise ValueError(f"Execution not found: {execution_id}")
|
||||
|
||||
return AgentExecutionStatus(
|
||||
execution_id=execution.id,
|
||||
workflow_id=execution.workflow_id,
|
||||
status=execution.status,
|
||||
current_step=execution.current_step,
|
||||
total_steps=execution.total_steps,
|
||||
step_states=execution.step_states,
|
||||
final_result=execution.final_result,
|
||||
error_message=execution.error_message,
|
||||
started_at=execution.started_at,
|
||||
completed_at=execution.completed_at,
|
||||
total_execution_time=execution.total_execution_time,
|
||||
total_cost=execution.total_cost,
|
||||
verification_proof=execution.verification_proof,
|
||||
)
|
||||
|
||||
async def _execute_steps_async(self, execution_id: str, inputs: dict[str, Any]) -> None:
|
||||
"""Execute workflow steps in dependency order"""
|
||||
|
||||
try:
|
||||
execution = await self.state_manager.get_execution(execution_id)
|
||||
workflow = await self.state_manager.get_workflow(execution.workflow_id)
|
||||
steps = await self.state_manager.get_workflow_steps(workflow.id)
|
||||
|
||||
# Build execution DAG
|
||||
step_order = self._build_execution_order(steps, workflow.dependencies)
|
||||
|
||||
current_inputs = inputs.copy()
|
||||
step_results = {}
|
||||
|
||||
for step_id in step_order:
|
||||
step = next(s for s in steps if s.id == step_id)
|
||||
|
||||
# Execute step
|
||||
step_result = await self._execute_single_step(execution_id, step, current_inputs)
|
||||
|
||||
step_results[step_id] = step_result
|
||||
|
||||
# Update inputs for next steps
|
||||
if step_result.output_data:
|
||||
current_inputs.update(step_result.output_data)
|
||||
|
||||
# Update execution progress
|
||||
await self.state_manager.update_execution_status(
|
||||
execution_id,
|
||||
current_step=execution.current_step + 1,
|
||||
completed_steps=execution.completed_steps + 1,
|
||||
step_states=step_results,
|
||||
)
|
||||
|
||||
# Mark execution as completed
|
||||
await self._complete_execution(execution_id, step_results)
|
||||
|
||||
except Exception as e:
|
||||
await self._handle_execution_failure(execution_id, e)
|
||||
|
||||
async def _execute_single_step(self, execution_id: str, step: AgentStep, inputs: dict[str, Any]) -> AgentStepExecution:
|
||||
"""Execute a single step"""
|
||||
|
||||
# Create step execution record
|
||||
step_execution = await self.state_manager.create_step_execution(execution_id, step.id)
|
||||
|
||||
try:
|
||||
# Update step status to running
|
||||
await self.state_manager.update_step_execution(
|
||||
step_execution.id, status=AgentStatus.RUNNING, started_at=datetime.now(timezone.utc), input_data=inputs
|
||||
)
|
||||
|
||||
# Execute the step based on type
|
||||
if step.step_type == StepType.INFERENCE:
|
||||
result = await self._execute_inference_step(step, inputs)
|
||||
elif step.step_type == StepType.TRAINING:
|
||||
result = await self._execute_training_step(step, inputs)
|
||||
elif step.step_type == StepType.DATA_PROCESSING:
|
||||
result = await self._execute_data_processing_step(step, inputs)
|
||||
else:
|
||||
result = await self._execute_custom_step(step, inputs)
|
||||
|
||||
# Update step execution with results
|
||||
await self.state_manager.update_step_execution(
|
||||
step_execution.id,
|
||||
status=AgentStatus.COMPLETED,
|
||||
completed_at=datetime.now(timezone.utc),
|
||||
output_data=result.get("output"),
|
||||
execution_time=result.get("execution_time", 0.0),
|
||||
gpu_accelerated=result.get("gpu_accelerated", False),
|
||||
memory_usage=result.get("memory_usage"),
|
||||
)
|
||||
|
||||
# Verify step if required
|
||||
if step.requires_proof:
|
||||
verification_result = await self.verifier.verify_step_execution(step_execution, step.verification_level)
|
||||
|
||||
await self.state_manager.update_step_execution(
|
||||
step_execution.id,
|
||||
step_proof=verification_result,
|
||||
verification_status="verified" if verification_result["verified"] else "failed",
|
||||
)
|
||||
|
||||
return step_execution
|
||||
|
||||
except Exception as e:
|
||||
# Mark step as failed
|
||||
await self.state_manager.update_step_execution(
|
||||
step_execution.id, status=AgentStatus.FAILED, completed_at=datetime.now(timezone.utc), error_message=str(e)
|
||||
)
|
||||
raise
|
||||
|
||||
async def _execute_inference_step(self, step: AgentStep, inputs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Execute inference step
|
||||
|
||||
Note: ML inference service integration requires:
|
||||
1. Connection to inference service (Ollama, custom API, etc.)
|
||||
2. Model selection and loading
|
||||
3. Input preprocessing and validation
|
||||
4. Output postprocessing
|
||||
Currently using simulated inference for testing purposes.
|
||||
"""
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
# Simulate processing time
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
execution_time = (datetime.now(timezone.utc) - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"output": {"prediction": "simulated_result", "confidence": 0.95},
|
||||
"execution_time": execution_time,
|
||||
"gpu_accelerated": False,
|
||||
"memory_usage": 128.5,
|
||||
}
|
||||
|
||||
async def _execute_training_step(self, step: AgentStep, inputs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Execute training step
|
||||
|
||||
Note: ML training service integration requires:
|
||||
1. Connection to training infrastructure (GPU clusters, distributed training)
|
||||
2. Dataset loading and preprocessing
|
||||
3. Training loop execution with monitoring
|
||||
4. Model checkpointing and validation
|
||||
Currently using simulated training for testing purposes.
|
||||
"""
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
# Simulate training time
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
execution_time = (datetime.now(timezone.utc) - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"output": {"model_updated": True, "training_loss": 0.123},
|
||||
"execution_time": execution_time,
|
||||
"gpu_accelerated": True, # Training typically uses GPU
|
||||
"memory_usage": 512.0,
|
||||
}
|
||||
|
||||
async def _execute_data_processing_step(self, step: AgentStep, inputs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Execute data processing step"""
|
||||
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
# Simulate processing time
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
execution_time = (datetime.now(timezone.utc) - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"output": {"processed_records": 1000, "data_validated": True},
|
||||
"execution_time": execution_time,
|
||||
"gpu_accelerated": False,
|
||||
"memory_usage": 64.0,
|
||||
}
|
||||
|
||||
async def _execute_custom_step(self, step: AgentStep, inputs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Execute custom step"""
|
||||
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
# Simulate custom processing
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
execution_time = (datetime.now(timezone.utc) - start_time).total_seconds()
|
||||
|
||||
return {
|
||||
"output": {"custom_result": "completed", "metadata": inputs},
|
||||
"execution_time": execution_time,
|
||||
"gpu_accelerated": False,
|
||||
"memory_usage": 256.0,
|
||||
}
|
||||
|
||||
def _build_execution_order(self, steps: list[AgentStep], dependencies: dict[str, list[str]]) -> list[str]:
|
||||
"""Build execution order based on dependencies"""
|
||||
|
||||
# Simple topological sort
|
||||
step_ids = [step.id for step in steps]
|
||||
ordered_steps = []
|
||||
remaining_steps = step_ids.copy()
|
||||
|
||||
while remaining_steps:
|
||||
# Find steps with no unmet dependencies
|
||||
ready_steps = []
|
||||
for step_id in remaining_steps:
|
||||
step_deps = dependencies.get(step_id, [])
|
||||
if all(dep in ordered_steps for dep in step_deps):
|
||||
ready_steps.append(step_id)
|
||||
|
||||
if not ready_steps:
|
||||
raise ValueError("Circular dependency detected in workflow")
|
||||
|
||||
# Add ready steps to order
|
||||
for step_id in ready_steps:
|
||||
ordered_steps.append(step_id)
|
||||
remaining_steps.remove(step_id)
|
||||
|
||||
return ordered_steps
|
||||
|
||||
async def _complete_execution(self, execution_id: str, step_results: dict[str, Any]) -> None:
|
||||
"""Mark execution as completed"""
|
||||
|
||||
completed_at = datetime.now(timezone.utc)
|
||||
execution = await self.state_manager.get_execution(execution_id)
|
||||
|
||||
total_execution_time = (completed_at - execution.started_at).total_seconds() if execution.started_at else 0.0
|
||||
|
||||
await self.state_manager.update_execution_status(
|
||||
execution_id,
|
||||
status=AgentStatus.COMPLETED,
|
||||
completed_at=completed_at,
|
||||
total_execution_time=total_execution_time,
|
||||
final_result={"step_results": step_results},
|
||||
)
|
||||
|
||||
async def _handle_execution_failure(self, execution_id: str, error: Exception) -> None:
|
||||
"""Handle execution failure"""
|
||||
|
||||
await self.state_manager.update_execution_status(
|
||||
execution_id, status=AgentStatus.FAILED, completed_at=datetime.now(timezone.utc), error_message=str(error)
|
||||
)
|
||||
|
||||
def _estimate_completion(self, execution: AgentExecution) -> datetime | None:
|
||||
"""Estimate completion time"""
|
||||
|
||||
if not execution.started_at:
|
||||
return None
|
||||
|
||||
# Simple estimation: 30 seconds per step
|
||||
estimated_duration = execution.total_steps * 30
|
||||
return execution.started_at + timedelta(seconds=estimated_duration)
|
||||
|
||||
def _estimate_cost(self, workflow: AIAgentWorkflow) -> float | None:
|
||||
"""Estimate total execution cost"""
|
||||
|
||||
# Simple cost model: $0.01 per step + base cost
|
||||
base_cost = 0.01
|
||||
per_step_cost = 0.01
|
||||
return base_cost + (len(workflow.steps) * per_step_cost)
|
||||
904
apps/agent-management/src/app/services/agent_service_marketplace.py
Executable file
904
apps/agent-management/src/app/services/agent_service_marketplace.py
Executable file
@@ -0,0 +1,904 @@
|
||||
"""
|
||||
AI Agent Service Marketplace Service
|
||||
Implements a sophisticated marketplace where agents can offer specialized services
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ServiceStatus(StrEnum):
|
||||
"""Service status types"""
|
||||
|
||||
ACTIVE = "active"
|
||||
INACTIVE = "inactive"
|
||||
SUSPENDED = "suspended"
|
||||
PENDING = "pending"
|
||||
|
||||
|
||||
class RequestStatus(StrEnum):
|
||||
"""Service request status types"""
|
||||
|
||||
PENDING = "pending"
|
||||
ACCEPTED = "accepted"
|
||||
COMPLETED = "completed"
|
||||
CANCELLED = "cancelled"
|
||||
EXPIRED = "expired"
|
||||
|
||||
|
||||
class GuildStatus(StrEnum):
|
||||
"""Guild status types"""
|
||||
|
||||
ACTIVE = "active"
|
||||
INACTIVE = "inactive"
|
||||
SUSPENDED = "suspended"
|
||||
|
||||
|
||||
class ServiceType(StrEnum):
|
||||
"""Service categories"""
|
||||
|
||||
DATA_ANALYSIS = "data_analysis"
|
||||
CONTENT_CREATION = "content_creation"
|
||||
RESEARCH = "research"
|
||||
CONSULTING = "consulting"
|
||||
DEVELOPMENT = "development"
|
||||
DESIGN = "design"
|
||||
MARKETING = "marketing"
|
||||
TRANSLATION = "translation"
|
||||
WRITING = "writing"
|
||||
ANALYSIS = "analysis"
|
||||
PREDICTION = "prediction"
|
||||
OPTIMIZATION = "optimization"
|
||||
AUTOMATION = "automation"
|
||||
MONITORING = "monitoring"
|
||||
TESTING = "testing"
|
||||
SECURITY = "security"
|
||||
INTEGRATION = "integration"
|
||||
CUSTOMIZATION = "customization"
|
||||
TRAINING = "training"
|
||||
SUPPORT = "support"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Service:
|
||||
"""Agent service information"""
|
||||
|
||||
id: str
|
||||
agent_id: str
|
||||
service_type: ServiceType
|
||||
name: str
|
||||
description: str
|
||||
metadata: dict[str, Any]
|
||||
base_price: float
|
||||
reputation: int
|
||||
status: ServiceStatus
|
||||
total_earnings: float
|
||||
completed_jobs: int
|
||||
average_rating: float
|
||||
rating_count: int
|
||||
listed_at: datetime
|
||||
last_updated: datetime
|
||||
guild_id: str | None = None
|
||||
tags: list[str] = field(default_factory=list)
|
||||
capabilities: list[str] = field(default_factory=list)
|
||||
requirements: list[str] = field(default_factory=list)
|
||||
pricing_model: str = "fixed" # fixed, hourly, per_task
|
||||
estimated_duration: int = 0 # in hours
|
||||
availability: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServiceRequest:
|
||||
"""Service request information"""
|
||||
|
||||
id: str
|
||||
client_id: str
|
||||
service_id: str
|
||||
budget: float
|
||||
requirements: str
|
||||
deadline: datetime
|
||||
status: RequestStatus
|
||||
assigned_agent: str | None = None
|
||||
accepted_at: datetime | None = None
|
||||
completed_at: datetime | None = None
|
||||
payment: float = 0.0
|
||||
rating: int = 0
|
||||
review: str = ""
|
||||
created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
results_hash: str | None = None
|
||||
priority: str = "normal" # low, normal, high, urgent
|
||||
complexity: str = "medium" # simple, medium, complex
|
||||
confidentiality: str = "public" # public, private, confidential
|
||||
|
||||
|
||||
@dataclass
|
||||
class Guild:
|
||||
"""Agent guild information"""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
description: str
|
||||
founder: str
|
||||
service_category: ServiceType
|
||||
member_count: int
|
||||
total_services: int
|
||||
total_earnings: float
|
||||
reputation: int
|
||||
status: GuildStatus
|
||||
created_at: datetime
|
||||
members: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||
requirements: list[str] = field(default_factory=list)
|
||||
benefits: list[str] = field(default_factory=list)
|
||||
guild_rules: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServiceCategory:
|
||||
"""Service category information"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
service_count: int
|
||||
total_volume: float
|
||||
average_price: float
|
||||
is_active: bool
|
||||
trending: bool = False
|
||||
popular_services: list[str] = field(default_factory=list)
|
||||
requirements: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MarketplaceAnalytics:
|
||||
"""Marketplace analytics data"""
|
||||
|
||||
total_services: int
|
||||
active_services: int
|
||||
total_requests: int
|
||||
pending_requests: int
|
||||
total_volume: float
|
||||
total_guilds: int
|
||||
average_service_price: float
|
||||
popular_categories: list[str]
|
||||
top_agents: list[str]
|
||||
revenue_trends: dict[str, float]
|
||||
growth_metrics: dict[str, float]
|
||||
|
||||
|
||||
class AgentServiceMarketplace:
|
||||
"""Service for managing AI agent service marketplace"""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
self.config = config
|
||||
self.services: dict[str, Service] = {}
|
||||
self.service_requests: dict[str, ServiceRequest] = {}
|
||||
self.guilds: dict[str, Guild] = {}
|
||||
self.categories: dict[str, ServiceCategory] = {}
|
||||
self.agent_services: dict[str, list[str]] = {}
|
||||
self.client_requests: dict[str, list[str]] = {}
|
||||
self.guild_services: dict[str, list[str]] = {}
|
||||
self.agent_guilds: dict[str, str] = {}
|
||||
self.services_by_type: dict[str, list[str]] = {}
|
||||
self.guilds_by_category: dict[str, list[str]] = {}
|
||||
|
||||
# Configuration
|
||||
self.marketplace_fee = 0.025 # 2.5%
|
||||
self.min_service_price = 0.001
|
||||
self.max_service_price = 1000.0
|
||||
self.min_reputation_to_list = 500
|
||||
self.request_timeout = 7 * 24 * 3600 # 7 days
|
||||
self.rating_weight = 100
|
||||
|
||||
# Initialize categories
|
||||
self._initialize_categories()
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize the marketplace service"""
|
||||
logger.info("Initializing Agent Service Marketplace")
|
||||
|
||||
# Load existing data
|
||||
await self._load_marketplace_data()
|
||||
|
||||
# Start background tasks
|
||||
asyncio.create_task(self._monitor_request_timeouts())
|
||||
asyncio.create_task(self._update_marketplace_analytics())
|
||||
asyncio.create_task(self._process_service_recommendations())
|
||||
asyncio.create_task(self._maintain_guild_reputation())
|
||||
|
||||
logger.info("Agent Service Marketplace initialized")
|
||||
|
||||
async def list_service(
|
||||
self,
|
||||
agent_id: str,
|
||||
service_type: ServiceType,
|
||||
name: str,
|
||||
description: str,
|
||||
metadata: dict[str, Any],
|
||||
base_price: float,
|
||||
tags: list[str],
|
||||
capabilities: list[str],
|
||||
requirements: list[str],
|
||||
pricing_model: str = "fixed",
|
||||
estimated_duration: int = 0,
|
||||
) -> Service:
|
||||
"""List a new service on the marketplace"""
|
||||
|
||||
try:
|
||||
# Validate inputs
|
||||
if base_price < self.min_service_price:
|
||||
raise ValueError(f"Price below minimum: {self.min_service_price}")
|
||||
|
||||
if base_price > self.max_service_price:
|
||||
raise ValueError(f"Price above maximum: {self.max_service_price}")
|
||||
|
||||
if not description or len(description) < 10:
|
||||
raise ValueError("Description too short")
|
||||
|
||||
# Check agent reputation (simplified - in production, check with reputation service)
|
||||
agent_reputation = await self._get_agent_reputation(agent_id)
|
||||
if agent_reputation < self.min_reputation_to_list:
|
||||
raise ValueError(f"Insufficient reputation: {agent_reputation}")
|
||||
|
||||
# Generate service ID
|
||||
service_id = await self._generate_service_id()
|
||||
|
||||
# Create service
|
||||
service = Service(
|
||||
id=service_id,
|
||||
agent_id=agent_id,
|
||||
service_type=service_type,
|
||||
name=name,
|
||||
description=description,
|
||||
metadata=metadata,
|
||||
base_price=base_price,
|
||||
reputation=agent_reputation,
|
||||
status=ServiceStatus.ACTIVE,
|
||||
total_earnings=0.0,
|
||||
completed_jobs=0,
|
||||
average_rating=0.0,
|
||||
rating_count=0,
|
||||
listed_at=datetime.now(timezone.utc),
|
||||
last_updated=datetime.now(timezone.utc),
|
||||
tags=tags,
|
||||
capabilities=capabilities,
|
||||
requirements=requirements,
|
||||
pricing_model=pricing_model,
|
||||
estimated_duration=estimated_duration,
|
||||
availability={
|
||||
"monday": True,
|
||||
"tuesday": True,
|
||||
"wednesday": True,
|
||||
"thursday": True,
|
||||
"friday": True,
|
||||
"saturday": False,
|
||||
"sunday": False,
|
||||
},
|
||||
)
|
||||
|
||||
# Store service
|
||||
self.services[service_id] = service
|
||||
|
||||
# Update mappings
|
||||
if agent_id not in self.agent_services:
|
||||
self.agent_services[agent_id] = []
|
||||
self.agent_services[agent_id].append(service_id)
|
||||
|
||||
if service_type.value not in self.services_by_type:
|
||||
self.services_by_type[service_type.value] = []
|
||||
self.services_by_type[service_type.value].append(service_id)
|
||||
|
||||
# Update category
|
||||
if service_type.value in self.categories:
|
||||
self.categories[service_type.value].service_count += 1
|
||||
|
||||
logger.info(f"Service listed: {service_id} by agent {agent_id}")
|
||||
return service
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list service: {e}")
|
||||
raise
|
||||
|
||||
async def request_service(
|
||||
self,
|
||||
client_id: str,
|
||||
service_id: str,
|
||||
budget: float,
|
||||
requirements: str,
|
||||
deadline: datetime,
|
||||
priority: str = "normal",
|
||||
complexity: str = "medium",
|
||||
confidentiality: str = "public",
|
||||
) -> ServiceRequest:
|
||||
"""Request a service"""
|
||||
|
||||
try:
|
||||
# Validate service
|
||||
if service_id not in self.services:
|
||||
raise ValueError(f"Service not found: {service_id}")
|
||||
|
||||
service = self.services[service_id]
|
||||
|
||||
if service.status != ServiceStatus.ACTIVE:
|
||||
raise ValueError("Service not active")
|
||||
|
||||
if budget < service.base_price:
|
||||
raise ValueError(f"Budget below service price: {service.base_price}")
|
||||
|
||||
if deadline <= datetime.now(timezone.utc):
|
||||
raise ValueError("Invalid deadline")
|
||||
|
||||
if deadline > datetime.now(timezone.utc) + timedelta(days=365):
|
||||
raise ValueError("Deadline too far in future")
|
||||
|
||||
# Generate request ID
|
||||
request_id = await self._generate_request_id()
|
||||
|
||||
# Create request
|
||||
request = ServiceRequest(
|
||||
id=request_id,
|
||||
client_id=client_id,
|
||||
service_id=service_id,
|
||||
budget=budget,
|
||||
requirements=requirements,
|
||||
deadline=deadline,
|
||||
status=RequestStatus.PENDING,
|
||||
priority=priority,
|
||||
complexity=complexity,
|
||||
confidentiality=confidentiality,
|
||||
)
|
||||
|
||||
# Store request
|
||||
self.service_requests[request_id] = request
|
||||
|
||||
# Update mappings
|
||||
if client_id not in self.client_requests:
|
||||
self.client_requests[client_id] = []
|
||||
self.client_requests[client_id].append(request_id)
|
||||
|
||||
# In production, transfer payment to escrow
|
||||
logger.info(f"Service requested: {request_id} for service {service_id}")
|
||||
return request
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to request service: {e}")
|
||||
raise
|
||||
|
||||
async def accept_request(self, request_id: str, agent_id: str) -> bool:
|
||||
"""Accept a service request"""
|
||||
|
||||
try:
|
||||
if request_id not in self.service_requests:
|
||||
raise ValueError(f"Request not found: {request_id}")
|
||||
|
||||
request = self.service_requests[request_id]
|
||||
service = self.services[request.service_id]
|
||||
|
||||
if request.status != RequestStatus.PENDING:
|
||||
raise ValueError("Request not pending")
|
||||
|
||||
if request.assigned_agent:
|
||||
raise ValueError("Request already assigned")
|
||||
|
||||
if service.agent_id != agent_id:
|
||||
raise ValueError("Not service provider")
|
||||
|
||||
if datetime.now(timezone.utc) > request.deadline:
|
||||
raise ValueError("Request expired")
|
||||
|
||||
# Update request
|
||||
request.status = RequestStatus.ACCEPTED
|
||||
request.assigned_agent = agent_id
|
||||
request.accepted_at = datetime.now(timezone.utc)
|
||||
|
||||
# Calculate dynamic price
|
||||
final_price = await self._calculate_dynamic_price(request.service_id, request.budget)
|
||||
request.payment = final_price
|
||||
|
||||
logger.info(f"Request accepted: {request_id} by agent {agent_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to accept request: {e}")
|
||||
raise
|
||||
|
||||
async def complete_request(self, request_id: str, agent_id: str, results: dict[str, Any]) -> bool:
|
||||
"""Complete a service request"""
|
||||
|
||||
try:
|
||||
if request_id not in self.service_requests:
|
||||
raise ValueError(f"Request not found: {request_id}")
|
||||
|
||||
request = self.service_requests[request_id]
|
||||
service = self.services[request.service_id]
|
||||
|
||||
if request.status != RequestStatus.ACCEPTED:
|
||||
raise ValueError("Request not accepted")
|
||||
|
||||
if request.assigned_agent != agent_id:
|
||||
raise ValueError("Not assigned agent")
|
||||
|
||||
if datetime.now(timezone.utc) > request.deadline:
|
||||
raise ValueError("Request expired")
|
||||
|
||||
# Update request
|
||||
request.status = RequestStatus.COMPLETED
|
||||
request.completed_at = datetime.now(timezone.utc)
|
||||
request.results_hash = hashlib.sha256(json.dumps(results, sort_keys=True).encode()).hexdigest()
|
||||
|
||||
# Calculate payment
|
||||
payment = request.payment
|
||||
fee = payment * self.marketplace_fee
|
||||
agent_payment = payment - fee
|
||||
|
||||
# Update service stats
|
||||
service.total_earnings += agent_payment
|
||||
service.completed_jobs += 1
|
||||
service.last_updated = datetime.now(timezone.utc)
|
||||
|
||||
# Update category
|
||||
if service.service_type.value in self.categories:
|
||||
self.categories[service.service_type.value].total_volume += payment
|
||||
|
||||
# Update guild stats
|
||||
if service.guild_id and service.guild_id in self.guilds:
|
||||
guild = self.guilds[service.guild_id]
|
||||
guild.total_earnings += agent_payment
|
||||
|
||||
# In production, process payment transfers
|
||||
logger.info(f"Request completed: {request_id} with payment {agent_payment}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to complete request: {e}")
|
||||
raise
|
||||
|
||||
async def rate_service(self, request_id: str, client_id: str, rating: int, review: str) -> bool:
|
||||
"""Rate and review a completed service"""
|
||||
|
||||
try:
|
||||
if request_id not in self.service_requests:
|
||||
raise ValueError(f"Request not found: {request_id}")
|
||||
|
||||
request = self.service_requests[request_id]
|
||||
service = self.services[request.service_id]
|
||||
|
||||
if request.status != RequestStatus.COMPLETED:
|
||||
raise ValueError("Request not completed")
|
||||
|
||||
if request.client_id != client_id:
|
||||
raise ValueError("Not request client")
|
||||
|
||||
if rating < 1 or rating > 5:
|
||||
raise ValueError("Invalid rating")
|
||||
|
||||
if datetime.now(timezone.utc) > request.deadline + timedelta(days=30):
|
||||
raise ValueError("Rating period expired")
|
||||
|
||||
# Update request
|
||||
request.rating = rating
|
||||
request.review = review
|
||||
|
||||
# Update service rating
|
||||
total_rating = service.average_rating * service.rating_count + rating
|
||||
service.rating_count += 1
|
||||
service.average_rating = total_rating / service.rating_count
|
||||
|
||||
# Update agent reputation
|
||||
reputation_change = await self._calculate_reputation_change(rating, service.reputation)
|
||||
await self._update_agent_reputation(service.agent_id, reputation_change)
|
||||
|
||||
logger.info(f"Service rated: {request_id} with rating {rating}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to rate service: {e}")
|
||||
raise
|
||||
|
||||
async def create_guild(
|
||||
self,
|
||||
founder_id: str,
|
||||
name: str,
|
||||
description: str,
|
||||
service_category: ServiceType,
|
||||
requirements: list[str],
|
||||
benefits: list[str],
|
||||
guild_rules: dict[str, Any],
|
||||
) -> Guild:
|
||||
"""Create a new guild"""
|
||||
|
||||
try:
|
||||
if not name or len(name) < 3:
|
||||
raise ValueError("Invalid guild name")
|
||||
|
||||
if service_category not in list(ServiceType):
|
||||
raise ValueError("Invalid service category")
|
||||
|
||||
# Generate guild ID
|
||||
guild_id = await self._generate_guild_id()
|
||||
|
||||
# Get founder reputation
|
||||
founder_reputation = await self._get_agent_reputation(founder_id)
|
||||
|
||||
# Create guild
|
||||
guild = Guild(
|
||||
id=guild_id,
|
||||
name=name,
|
||||
description=description,
|
||||
founder=founder_id,
|
||||
service_category=service_category,
|
||||
member_count=1,
|
||||
total_services=0,
|
||||
total_earnings=0.0,
|
||||
reputation=founder_reputation,
|
||||
status=GuildStatus.ACTIVE,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
requirements=requirements,
|
||||
benefits=benefits,
|
||||
guild_rules=guild_rules,
|
||||
)
|
||||
|
||||
# Add founder as member
|
||||
guild.members[founder_id] = {
|
||||
"joined_at": datetime.now(timezone.utc),
|
||||
"reputation": founder_reputation,
|
||||
"role": "founder",
|
||||
"contributions": 0,
|
||||
}
|
||||
|
||||
# Store guild
|
||||
self.guilds[guild_id] = guild
|
||||
|
||||
# Update mappings
|
||||
if service_category.value not in self.guilds_by_category:
|
||||
self.guilds_by_category[service_category.value] = []
|
||||
self.guilds_by_category[service_category.value].append(guild_id)
|
||||
|
||||
self.agent_guilds[founder_id] = guild_id
|
||||
|
||||
logger.info(f"Guild created: {guild_id} by {founder_id}")
|
||||
return guild
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create guild: {e}")
|
||||
raise
|
||||
|
||||
async def join_guild(self, agent_id: str, guild_id: str) -> bool:
|
||||
"""Join a guild"""
|
||||
|
||||
try:
|
||||
if guild_id not in self.guilds:
|
||||
raise ValueError(f"Guild not found: {guild_id}")
|
||||
|
||||
guild = self.guilds[guild_id]
|
||||
|
||||
if agent_id in guild.members:
|
||||
raise ValueError("Already a member")
|
||||
|
||||
if guild.status != GuildStatus.ACTIVE:
|
||||
raise ValueError("Guild not active")
|
||||
|
||||
# Check agent reputation
|
||||
agent_reputation = await self._get_agent_reputation(agent_id)
|
||||
if agent_reputation < guild.reputation // 2:
|
||||
raise ValueError("Insufficient reputation")
|
||||
|
||||
# Add member
|
||||
guild.members[agent_id] = {
|
||||
"joined_at": datetime.now(timezone.utc),
|
||||
"reputation": agent_reputation,
|
||||
"role": "member",
|
||||
"contributions": 0,
|
||||
}
|
||||
guild.member_count += 1
|
||||
|
||||
# Update mappings
|
||||
self.agent_guilds[agent_id] = guild_id
|
||||
|
||||
logger.info(f"Agent {agent_id} joined guild {guild_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to join guild: {e}")
|
||||
raise
|
||||
|
||||
async def search_services(
|
||||
self,
|
||||
query: str | None = None,
|
||||
service_type: ServiceType | None = None,
|
||||
tags: list[str] | None = None,
|
||||
min_price: float | None = None,
|
||||
max_price: float | None = None,
|
||||
min_rating: float | None = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
) -> list[Service]:
|
||||
"""Search services with various filters"""
|
||||
|
||||
try:
|
||||
results = []
|
||||
|
||||
# Filter through all services
|
||||
for service in self.services.values():
|
||||
if service.status != ServiceStatus.ACTIVE:
|
||||
continue
|
||||
|
||||
# Apply filters
|
||||
if service_type and service.service_type != service_type:
|
||||
continue
|
||||
|
||||
if min_price and service.base_price < min_price:
|
||||
continue
|
||||
|
||||
if max_price and service.base_price > max_price:
|
||||
continue
|
||||
|
||||
if min_rating and service.average_rating < min_rating:
|
||||
continue
|
||||
|
||||
if tags and not any(tag in service.tags for tag in tags):
|
||||
continue
|
||||
|
||||
if query:
|
||||
query_lower = query.lower()
|
||||
if (
|
||||
query_lower not in service.name.lower()
|
||||
and query_lower not in service.description.lower()
|
||||
and not any(query_lower in tag.lower() for tag in service.tags)
|
||||
):
|
||||
continue
|
||||
|
||||
results.append(service)
|
||||
|
||||
# Sort by relevance (simplified)
|
||||
results.sort(key=lambda x: (x.average_rating, x.reputation), reverse=True)
|
||||
|
||||
# Apply pagination
|
||||
return results[offset : offset + limit]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to search services: {e}")
|
||||
raise
|
||||
|
||||
async def get_agent_services(self, agent_id: str) -> list[Service]:
|
||||
"""Get all services for an agent"""
|
||||
|
||||
try:
|
||||
if agent_id not in self.agent_services:
|
||||
return []
|
||||
|
||||
services = []
|
||||
for service_id in self.agent_services[agent_id]:
|
||||
if service_id in self.services:
|
||||
services.append(self.services[service_id])
|
||||
|
||||
return services
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get agent services: {e}")
|
||||
raise
|
||||
|
||||
async def get_client_requests(self, client_id: str) -> list[ServiceRequest]:
|
||||
"""Get all requests for a client"""
|
||||
|
||||
try:
|
||||
if client_id not in self.client_requests:
|
||||
return []
|
||||
|
||||
requests = []
|
||||
for request_id in self.client_requests[client_id]:
|
||||
if request_id in self.service_requests:
|
||||
requests.append(self.service_requests[request_id])
|
||||
|
||||
return requests
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get client requests: {e}")
|
||||
raise
|
||||
|
||||
async def get_marketplace_analytics(self) -> MarketplaceAnalytics:
|
||||
"""Get marketplace analytics"""
|
||||
|
||||
try:
|
||||
total_services = len(self.services)
|
||||
active_services = len([s for s in self.services.values() if s.status == ServiceStatus.ACTIVE])
|
||||
total_requests = len(self.service_requests)
|
||||
pending_requests = len([r for r in self.service_requests.values() if r.status == RequestStatus.PENDING])
|
||||
total_guilds = len(self.guilds)
|
||||
|
||||
# Calculate total volume
|
||||
total_volume = sum(service.total_earnings for service in self.services.values())
|
||||
|
||||
# Calculate average price
|
||||
active_service_prices = [
|
||||
service.base_price for service in self.services.values() if service.status == ServiceStatus.ACTIVE
|
||||
]
|
||||
average_price = sum(active_service_prices) / len(active_service_prices) if active_service_prices else 0
|
||||
|
||||
# Get popular categories
|
||||
category_counts = {}
|
||||
for service in self.services.values():
|
||||
if service.status == ServiceStatus.ACTIVE:
|
||||
category_counts[service.service_type.value] = category_counts.get(service.service_type.value, 0) + 1
|
||||
|
||||
popular_categories = sorted(category_counts.items(), key=lambda x: x[1], reverse=True)[:5]
|
||||
|
||||
# Get top agents
|
||||
agent_earnings = {}
|
||||
for service in self.services.values():
|
||||
agent_earnings[service.agent_id] = agent_earnings.get(service.agent_id, 0) + service.total_earnings
|
||||
|
||||
top_agents = sorted(agent_earnings.items(), key=lambda x: x[1], reverse=True)[:5]
|
||||
|
||||
return MarketplaceAnalytics(
|
||||
total_services=total_services,
|
||||
active_services=active_services,
|
||||
total_requests=total_requests,
|
||||
pending_requests=pending_requests,
|
||||
total_volume=total_volume,
|
||||
total_guilds=total_guilds,
|
||||
average_service_price=average_price,
|
||||
popular_categories=[cat[0] for cat in popular_categories],
|
||||
top_agents=[agent[0] for agent in top_agents],
|
||||
revenue_trends={},
|
||||
growth_metrics={},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get marketplace analytics: {e}")
|
||||
raise
|
||||
|
||||
async def _calculate_dynamic_price(self, service_id: str, budget: float) -> float:
|
||||
"""Calculate dynamic price based on demand and reputation"""
|
||||
|
||||
service = self.services[service_id]
|
||||
dynamic_price = service.base_price
|
||||
|
||||
# Reputation multiplier
|
||||
reputation_multiplier = 1.0 + (service.reputation / 10000) * 0.5
|
||||
dynamic_price *= reputation_multiplier
|
||||
|
||||
# Demand multiplier
|
||||
demand_multiplier = 1.0
|
||||
if service.completed_jobs > 10:
|
||||
demand_multiplier = 1.0 + (service.completed_jobs / 100) * 0.5
|
||||
dynamic_price *= demand_multiplier
|
||||
|
||||
# Rating multiplier
|
||||
rating_multiplier = 1.0 + (service.average_rating / 5) * 0.3
|
||||
dynamic_price *= rating_multiplier
|
||||
|
||||
return min(dynamic_price, budget)
|
||||
|
||||
async def _calculate_reputation_change(self, rating: int, current_reputation: int) -> int:
|
||||
"""Calculate reputation change based on rating"""
|
||||
|
||||
if rating == 5:
|
||||
return self.rating_weight * 2
|
||||
elif rating == 4:
|
||||
return self.rating_weight
|
||||
elif rating == 3:
|
||||
return 0
|
||||
elif rating == 2:
|
||||
return -self.rating_weight
|
||||
else: # rating == 1
|
||||
return -self.rating_weight * 2
|
||||
|
||||
async def _get_agent_reputation(self, agent_id: str) -> int:
|
||||
"""Get agent reputation (simplified)"""
|
||||
# In production, integrate with reputation service
|
||||
return 1000
|
||||
|
||||
async def _update_agent_reputation(self, agent_id: str, change: int):
|
||||
"""Update agent reputation (simplified)"""
|
||||
# In production, integrate with reputation service
|
||||
pass
|
||||
|
||||
async def _generate_service_id(self) -> str:
|
||||
"""Generate unique service ID"""
|
||||
import uuid
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
async def _generate_request_id(self) -> str:
|
||||
"""Generate unique request ID"""
|
||||
import uuid
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
async def _generate_guild_id(self) -> str:
|
||||
"""Generate unique guild ID"""
|
||||
import uuid
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def _initialize_categories(self):
|
||||
"""Initialize service categories"""
|
||||
|
||||
for service_type in ServiceType:
|
||||
self.categories[service_type.value] = ServiceCategory(
|
||||
name=service_type.value,
|
||||
description=f"Services related to {service_type.value}",
|
||||
service_count=0,
|
||||
total_volume=0.0,
|
||||
average_price=0.0,
|
||||
is_active=True,
|
||||
)
|
||||
|
||||
async def _load_marketplace_data(self):
|
||||
"""Load existing marketplace data"""
|
||||
# In production, load from database
|
||||
pass
|
||||
|
||||
async def _monitor_request_timeouts(self):
|
||||
"""Monitor and handle request timeouts"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
current_time = datetime.now(timezone.utc)
|
||||
|
||||
for request in self.service_requests.values():
|
||||
if request.status == RequestStatus.PENDING and current_time > request.deadline:
|
||||
request.status = RequestStatus.EXPIRED
|
||||
logger.info(f"Request expired: {request.id}")
|
||||
|
||||
await asyncio.sleep(3600) # Check every hour
|
||||
except Exception as e:
|
||||
logger.error(f"Error monitoring timeouts: {e}")
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
async def _update_marketplace_analytics(self):
|
||||
"""Update marketplace analytics"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Update trending categories
|
||||
for category in self.categories.values():
|
||||
# Simplified trending logic
|
||||
category.trending = category.service_count > 10
|
||||
|
||||
await asyncio.sleep(3600) # Update every hour
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating analytics: {e}")
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
async def _process_service_recommendations(self):
|
||||
"""Process service recommendations"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Implement recommendation logic
|
||||
await asyncio.sleep(1800) # Process every 30 minutes
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing recommendations: {e}")
|
||||
await asyncio.sleep(1800)
|
||||
|
||||
async def _maintain_guild_reputation(self):
|
||||
"""Maintain guild reputation scores"""
|
||||
|
||||
while True:
|
||||
try:
|
||||
for guild in self.guilds.values():
|
||||
# Calculate guild reputation based on members
|
||||
total_reputation = 0
|
||||
active_members = 0
|
||||
|
||||
for member_id, _member_data in guild.members.items():
|
||||
member_reputation = await self._get_agent_reputation(member_id)
|
||||
total_reputation += member_reputation
|
||||
active_members += 1
|
||||
|
||||
if active_members > 0:
|
||||
guild.reputation = total_reputation // active_members
|
||||
|
||||
await asyncio.sleep(3600) # Update every hour
|
||||
except Exception as e:
|
||||
logger.error(f"Error maintaining guild reputation: {e}")
|
||||
await asyncio.sleep(3600)
|
||||
@@ -19,6 +19,8 @@ logger = logging.getLogger(__name__)
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '../../../..'))
|
||||
|
||||
from apps.agent_services.agent_bridge.src.integration_layer import AgentServiceBridge
|
||||
from aitbc import get_logger
|
||||
logger = get_logger(__name__)
|
||||
|
||||
class ComplianceAgent:
|
||||
"""Automated compliance agent"""
|
||||
@@ -142,11 +144,11 @@ async def main():
|
||||
# Run compliance loop
|
||||
await agent.run_compliance_loop()
|
||||
except KeyboardInterrupt:
|
||||
print("Shutting down compliance agent...")
|
||||
logger.info("Shutting down compliance agent...")
|
||||
finally:
|
||||
await agent.stop()
|
||||
else:
|
||||
print("Failed to start compliance agent")
|
||||
logger.error("Failed to start compliance agent")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
@@ -15,6 +15,9 @@ import sqlite3
|
||||
from contextlib import contextmanager
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from aitbc import get_logger
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Use absolute path for database in /var/lib/aitbc for persistence
|
||||
DB_DIR = "/var/lib/aitbc"
|
||||
os.makedirs(DB_DIR, exist_ok=True)
|
||||
@@ -145,9 +148,9 @@ async def create_task(task: TaskCreation):
|
||||
assigned_agent_id = assign_task_to_agent(task_id, task.required_capabilities)
|
||||
|
||||
if assigned_agent_id:
|
||||
print(f"Task {task_id} assigned to agent {assigned_agent_id}")
|
||||
logger.info(f"Task {task_id} assigned to agent {assigned_agent_id}")
|
||||
else:
|
||||
print(f"Task {task_id} - no eligible agents found")
|
||||
logger.info(f"Task {task_id} - no eligible agents found")
|
||||
|
||||
return Task(
|
||||
id=task_id,
|
||||
@@ -193,7 +196,7 @@ async def health_check():
|
||||
@app.get("/tasks/status")
|
||||
async def get_task_status():
|
||||
"""Get task distribution statistics including active agents"""
|
||||
print(f"DEBUG: Querying tasks/status, DB_PATH={DB_PATH}")
|
||||
logger.debug(f"DEBUG: Querying tasks/status, DB_PATH={DB_PATH}")
|
||||
with get_db_connection() as conn:
|
||||
# Get task statistics
|
||||
tasks = conn.execute("SELECT * FROM tasks").fetchall()
|
||||
@@ -203,7 +206,7 @@ async def get_task_status():
|
||||
|
||||
# Get active agents count
|
||||
agents = conn.execute("SELECT * FROM agents WHERE status = ?", ("active",)).fetchall()
|
||||
print(f"DEBUG: Found {len(agents)} active agents")
|
||||
logger.debug(f"DEBUG: Found {len(agents)} active agents")
|
||||
active_agents = len(agents)
|
||||
|
||||
# Calculate load balancer stats
|
||||
@@ -256,11 +259,11 @@ async def get_task_status():
|
||||
async def register_agent(request: AgentRegistrationRequest):
|
||||
"""Register a new agent"""
|
||||
try:
|
||||
print(f"DEBUG: Attempting to register agent {request.agent_id}")
|
||||
print(f"DEBUG: Database path: {DB_PATH}")
|
||||
logger.debug(f"DEBUG: Attempting to register agent {request.agent_id}")
|
||||
logger.debug(f"DEBUG: Database path: {DB_PATH}")
|
||||
conn = get_db()
|
||||
try:
|
||||
print(f"DEBUG: Database connection established")
|
||||
logger.debug(f"DEBUG: Database connection established")
|
||||
conn.execute('''
|
||||
INSERT INTO agents (id, agent_type, status, capabilities, services, endpoints, metadata, last_heartbeat, health_score)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
@@ -276,7 +279,7 @@ async def register_agent(request: AgentRegistrationRequest):
|
||||
1.0
|
||||
))
|
||||
conn.commit()
|
||||
print(f"DEBUG: Agent {request.agent_id} inserted and committed")
|
||||
logger.debug(f"DEBUG: Agent {request.agent_id} inserted and committed")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@@ -287,7 +290,7 @@ async def register_agent(request: AgentRegistrationRequest):
|
||||
"registered_at": datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to register agent: {str(e)}")
|
||||
logger.error(f"ERROR: Failed to register agent: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=f"Failed to register agent: {str(e)}")
|
||||
|
||||
@app.post("/agents/discover")
|
||||
|
||||
@@ -16,6 +16,8 @@ import os
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '../../../..'))
|
||||
|
||||
from apps.agent_services.agent_bridge.src.integration_layer import AgentServiceBridge
|
||||
from aitbc import get_logger
|
||||
logger = get_logger(__name__)
|
||||
|
||||
class TradingAgent:
|
||||
"""Automated trading agent"""
|
||||
@@ -156,11 +158,11 @@ async def main():
|
||||
# Run trading loop
|
||||
await agent.run_trading_loop()
|
||||
except KeyboardInterrupt:
|
||||
print("Shutting down trading agent...")
|
||||
logger.info("Shutting down trading agent...")
|
||||
finally:
|
||||
await agent.stop()
|
||||
else:
|
||||
print("Failed to start trading agent")
|
||||
logger.error("Failed to start trading agent")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
0
apps/ai-models/src/app/__init__.py
Normal file
0
apps/ai-models/src/app/__init__.py
Normal file
0
apps/ai-models/src/app/models/__init__.py
Normal file
0
apps/ai-models/src/app/models/__init__.py
Normal file
0
apps/ai-models/src/app/routers/__init__.py
Normal file
0
apps/ai-models/src/app/routers/__init__.py
Normal file
0
apps/ai-models/src/app/services/__init__.py
Normal file
0
apps/ai-models/src/app/services/__init__.py
Normal file
@@ -9,7 +9,7 @@ python = "^3.13"
|
||||
fastapi = ">=0.115.6"
|
||||
uvicorn = "^0.24.0"
|
||||
httpx = ">=0.28.1"
|
||||
aitbc-core = {path = "../../packages/py/aitbc-core", develop = true}
|
||||
|
||||
|
||||
[tool.poetry.group.test.dependencies]
|
||||
pytest = ">=9.0.3"
|
||||
|
||||
@@ -12,6 +12,8 @@ from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
from cryptography.hazmat.primitives.serialization import Encoding, PrivateFormat, NoEncryption
|
||||
from aitbc import get_logger
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@dataclass
|
||||
class ValidatorKeyPair:
|
||||
@@ -52,7 +54,7 @@ class KeyManager:
|
||||
last_rotated=key_data['last_rotated']
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error loading keys: {e}")
|
||||
logger.error(f"Error loading keys: {e}")
|
||||
|
||||
def generate_key_pair(self, address: str) -> ValidatorKeyPair:
|
||||
"""Generate new RSA key pair for validator"""
|
||||
@@ -195,7 +197,7 @@ class KeyManager:
|
||||
# Set secure permissions
|
||||
os.chmod(keys_file, 0o600)
|
||||
except Exception as e:
|
||||
print(f"Error saving keys: {e}")
|
||||
logger.error("Error saving keys", error=str(e))
|
||||
|
||||
def should_rotate_key(self, address: str, rotation_interval: int = 86400) -> bool:
|
||||
"""Check if key should be rotated (default: 24 hours)"""
|
||||
|
||||
@@ -11,9 +11,18 @@ from dataclasses import dataclass, asdict
|
||||
from enum import Enum
|
||||
from decimal import Decimal
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
def log_info(message: str):
|
||||
"""Simple logging function"""
|
||||
print(f"[EscrowManager] {message}")
|
||||
logger.info(message)
|
||||
|
||||
# Remove the old print-based logging function
|
||||
def log_info_old(message: str):
|
||||
"""Legacy logging function - use logger instead"""
|
||||
logger.info(f"[EscrowManager] {message}")
|
||||
|
||||
class EscrowState(Enum):
|
||||
CREATED = "created"
|
||||
|
||||
@@ -12,6 +12,10 @@ from sqlalchemy.orm import sessionmaker, Session
|
||||
from eth_utils import to_checksum_address
|
||||
import json
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
@@ -168,7 +172,7 @@ class PersistentSpendingTracker:
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to record spending: {e}")
|
||||
logger.error(f"Failed to record spending: {e}")
|
||||
return False
|
||||
|
||||
def check_spending_limits(self, agent_address: str, amount: float, timestamp: datetime = None) -> SpendingCheckResult:
|
||||
@@ -332,7 +336,7 @@ class PersistentSpendingTracker:
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to update spending limits: {e}")
|
||||
logger.error("Failed to update spending limits", error=str(e))
|
||||
return False
|
||||
|
||||
def add_guardian(self, agent_address: str, guardian_address: str, added_by: str) -> bool:
|
||||
@@ -378,7 +382,7 @@ class PersistentSpendingTracker:
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to add guardian: {e}")
|
||||
logger.error("Failed to add guardian", error=str(e))
|
||||
return False
|
||||
|
||||
def is_guardian_authorized(self, agent_address: str, guardian_address: str) -> bool:
|
||||
|
||||
@@ -34,7 +34,7 @@ async def main() -> None:
|
||||
)
|
||||
try:
|
||||
imported = await sync.bulk_import_from(args.source, import_url=args.import_url)
|
||||
print(f"[+] Bulk sync complete: imported {imported} blocks")
|
||||
logger.info("Bulk sync complete", blocks_imported=imported)
|
||||
finally:
|
||||
await sync.close()
|
||||
|
||||
|
||||
0
apps/blockchain/src/app/__init__.py
Normal file
0
apps/blockchain/src/app/__init__.py
Normal file
0
apps/blockchain/src/app/models/__init__.py
Normal file
0
apps/blockchain/src/app/models/__init__.py
Normal file
0
apps/blockchain/src/app/routers/__init__.py
Normal file
0
apps/blockchain/src/app/routers/__init__.py
Normal file
0
apps/blockchain/src/app/services/__init__.py
Normal file
0
apps/blockchain/src/app/services/__init__.py
Normal file
0
apps/computing/src/app/__init__.py
Normal file
0
apps/computing/src/app/__init__.py
Normal file
0
apps/computing/src/app/models/__init__.py
Normal file
0
apps/computing/src/app/models/__init__.py
Normal file
0
apps/computing/src/app/routers/__init__.py
Normal file
0
apps/computing/src/app/routers/__init__.py
Normal file
0
apps/computing/src/app/services/__init__.py
Normal file
0
apps/computing/src/app/services/__init__.py
Normal file
@@ -5,7 +5,7 @@ DATABASE_URL=sqlite:////var/lib/aitbc/data/coordinator.db
|
||||
CLIENT_API_KEYS=${CLIENT_API_KEY},client_dev_key_2
|
||||
MINER_API_KEYS=${MINER_API_KEY},miner_dev_key_2
|
||||
ADMIN_API_KEYS=${ADMIN_API_KEY}
|
||||
HMAC_SECRET=change_me
|
||||
HMAC_SECRET=
|
||||
ALLOW_ORIGINS=*
|
||||
JOB_TTL_SECONDS=900
|
||||
HEARTBEAT_INTERVAL_SECONDS=10
|
||||
|
||||
126
apps/coordinator-api/DECOMPOSITION_PROGRESS.md
Normal file
126
apps/coordinator-api/DECOMPOSITION_PROGRESS.md
Normal file
@@ -0,0 +1,126 @@
|
||||
# Coordinator-API Decomposition Progress
|
||||
|
||||
## Phase 1: Modular Monolith Restructuring (Completed)
|
||||
|
||||
### Week 1: Domain Boundary Identification ✓
|
||||
|
||||
**Completed Tasks:**
|
||||
- Mapped 61 routers to bounded contexts
|
||||
- Identified cross-context dependencies between routers and services
|
||||
- Created context-specific subdirectory structure for:
|
||||
- `contexts/marketplace/` (routers, services, domain, storage)
|
||||
- `contexts/payments/` (routers, services, domain, storage)
|
||||
- `contexts/blockchain/` (routers, services, domain, storage)
|
||||
- `contexts/agent_identity/` (routers, services, domain, storage)
|
||||
|
||||
### Week 2: Service Layer Extraction ✓
|
||||
|
||||
**Completed Tasks:**
|
||||
- Extracted context-specific services to context directories:
|
||||
- Marketplace: marketplace.py, marketplace_enhanced.py, marketplace_enhanced_simple.py, global_marketplace.py, global_marketplace_integration.py
|
||||
- Payments: payments.py
|
||||
- Blockchain: blockchain.py
|
||||
- Agent Identity: (already existed in agent_identity/ directory)
|
||||
- Extracted domain models to context directories:
|
||||
- Marketplace: marketplace.py, gpu_marketplace.py, global_marketplace.py
|
||||
- Payments: payment.py
|
||||
- Agent Identity: agent_identity.py
|
||||
- Updated all imports in moved files to reference correct paths
|
||||
- Created __init__.py files for all context directories
|
||||
|
||||
### Week 3: Router Organization ✓
|
||||
|
||||
**Completed Tasks:**
|
||||
- Moved routers to context directories:
|
||||
- Marketplace: marketplace.py, marketplace_gpu.py, marketplace_offers.py, global_marketplace.py, global_marketplace_integration.py
|
||||
- Payments: payments.py
|
||||
- Blockchain: blockchain.py
|
||||
- Agent Identity: agent_identity.py
|
||||
- Updated main.py to register routers from new context locations
|
||||
- All imports updated to use context-qualified paths
|
||||
- Fixed pre-existing syntax error in governance.py
|
||||
|
||||
### Week 4: Database Schema Separation ✓
|
||||
|
||||
**Completed Tasks:**
|
||||
- Created context-specific SQLAlchemy schema files:
|
||||
- `contexts/marketplace/storage/schema.py` - defines marketplace_ prefix
|
||||
- `contexts/payments/storage/schema.py` - defines payments_ prefix
|
||||
- `contexts/blockchain/storage/schema.py` - defines blockchain_ prefix
|
||||
- `contexts/agent_identity/storage/schema.py` - defines agent_identity_ prefix
|
||||
- Updated domain models to use context-prefixed table names:
|
||||
- Marketplace: MarketplaceOffer -> marketplace_offer, MarketplaceBid -> marketplace_bid
|
||||
- Payments: JobPayment -> payments_job_payment, PaymentEscrow -> payments_escrow
|
||||
- Agent Identity: AgentIdentity -> agent_identity_identity, CrossChainMapping -> agent_identity_cross_chain_mapping, IdentityVerification -> agent_identity_verification
|
||||
- Created Alembic migration script: `alembic/versions/001_context_table_prefixes.py`
|
||||
- Compilation verified successfully after table name changes
|
||||
|
||||
## Current State
|
||||
|
||||
**Compilation Status:** ✓ PASSED
|
||||
- All Python files in coordinator-api compile successfully
|
||||
- No import errors after restructuring
|
||||
- main.py successfully imports routers from context directories
|
||||
|
||||
**Code Metrics:**
|
||||
- Contexts created: 4 (marketplace, payments, blockchain, agent_identity)
|
||||
- Routers moved: 8
|
||||
- Services moved: 8
|
||||
- Domain models moved: 5
|
||||
- Import paths updated: 21 files
|
||||
|
||||
## Next Steps (Phase 2: Microservice Extraction)
|
||||
|
||||
According to the decomposition plan, Phase 2 involves:
|
||||
1. Week 5: Marketplace Service Extraction
|
||||
2. Week 6: Agent Identity Service Extraction
|
||||
3. Week 7: Payments Service Extraction
|
||||
4. Week 8: Validation & Monitoring
|
||||
|
||||
## Files Modified
|
||||
|
||||
**Created:**
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/marketplace/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/marketplace/routers/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/marketplace/services/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/marketplace/domain/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/marketplace/storage/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/payments/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/payments/routers/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/payments/services/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/payments/domain/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/payments/storage/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/blockchain/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/blockchain/routers/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/blockchain/services/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/blockchain/domain/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/blockchain/storage/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/agent_identity/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/agent_identity/routers/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/agent_identity/services/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/agent_identity/domain/__init__.py`
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/contexts/agent_identity/storage/__init__.py`
|
||||
|
||||
**Modified:**
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/main.py` - Updated router imports
|
||||
- `/opt/aitbc/apps/coordinator-api/src/app/routers/governance.py` - Fixed syntax error
|
||||
|
||||
**Moved (Routers):**
|
||||
- marketplace.py, marketplace_gpu.py, marketplace_offers.py, global_marketplace.py, global_marketplace_integration.py → contexts/marketplace/routers/
|
||||
- payments.py → contexts/payments/routers/
|
||||
- blockchain.py → contexts/blockchain/routers/
|
||||
- agent_identity.py → contexts/agent_identity/routers/
|
||||
|
||||
**Moved (Services):**
|
||||
- marketplace.py, marketplace_enhanced.py, marketplace_enhanced_simple.py, global_marketplace.py, global_marketplace_integration.py → contexts/marketplace/services/
|
||||
- payments.py → contexts/payments/services/
|
||||
- blockchain.py → contexts/blockchain/services/
|
||||
|
||||
**Moved (Domain):**
|
||||
- marketplace.py, gpu_marketplace.py, global_marketplace.py → contexts/marketplace/domain/
|
||||
- payment.py → contexts/payments/domain/
|
||||
- agent_identity.py → contexts/agent_identity/domain/
|
||||
|
||||
**Import Updates:**
|
||||
- All moved files updated with correct relative import paths (e.g., `..` → `....` for routers, `..` → `....` for services)
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Add context prefixes to table names
|
||||
|
||||
Revision ID: 001_context_prefixes
|
||||
Revises:
|
||||
Create Date: 2026-05-12
|
||||
|
||||
This migration renames tables to use context-specific prefixes:
|
||||
- marketplaceoffer -> marketplace_offer
|
||||
- marketplacebid -> marketplace_bid
|
||||
- job_payments -> payments_job_payment
|
||||
- payment_escrows -> payments_escrow
|
||||
- agent_identities -> agent_identity_identity
|
||||
- cross_chain_mappings -> agent_identity_cross_chain_mapping
|
||||
- identity_verifications -> agent_identity_verification
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '001_context_prefixes'
|
||||
down_revision = None
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Marketplace context table renames
|
||||
op.rename_table('marketplaceoffer', 'marketplace_offer')
|
||||
op.rename_table('marketplacebid', 'marketplace_bid')
|
||||
|
||||
# Payments context table renames
|
||||
op.rename_table('job_payments', 'payments_job_payment')
|
||||
op.rename_table('payment_escrows', 'payments_escrow')
|
||||
|
||||
# Agent Identity context table renames
|
||||
op.rename_table('agent_identities', 'agent_identity_identity')
|
||||
op.rename_table('cross_chain_mappings', 'agent_identity_cross_chain_mapping')
|
||||
op.rename_table('identity_verifications', 'agent_identity_verification')
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Reverse the renames
|
||||
op.rename_table('marketplace_offer', 'marketplaceoffer')
|
||||
op.rename_table('marketplace_bid', 'marketplacebid')
|
||||
|
||||
op.rename_table('payments_job_payment', 'job_payments')
|
||||
op.rename_table('payments_escrow', 'payment_escrows')
|
||||
|
||||
op.rename_table('agent_identity_identity', 'agent_identities')
|
||||
op.rename_table('agent_identity_cross_chain_mapping', 'cross_chain_mappings')
|
||||
op.rename_table('agent_identity_verification', 'identity_verifications')
|
||||
@@ -50,12 +50,12 @@ def migrate_all_data():
|
||||
print(f" Skipping table {table_name} (not in allowed list)")
|
||||
continue
|
||||
|
||||
sqlite_cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
sqlite_cursor.execute(f"PRAGMA table_info(\"{table_name}\")")
|
||||
columns = sqlite_cursor.fetchall()
|
||||
column_names = [col[1] for col in columns]
|
||||
|
||||
# Get data
|
||||
sqlite_cursor.execute(f"SELECT * FROM {table_name}")
|
||||
sqlite_cursor.execute(f"SELECT * FROM \"{table_name}\"")
|
||||
rows = sqlite_cursor.fetchall()
|
||||
|
||||
if not rows:
|
||||
@@ -70,7 +70,7 @@ def migrate_all_data():
|
||||
'''
|
||||
else:
|
||||
insert_sql = f'''
|
||||
INSERT INTO {table_name} ({', '.join(column_names)})
|
||||
INSERT INTO "{table_name}" ({', '.join(column_names)})
|
||||
VALUES ({', '.join(['%s'] * len(column_names))})
|
||||
'''
|
||||
|
||||
|
||||
@@ -261,7 +261,7 @@ def migrate_data():
|
||||
continue
|
||||
|
||||
print(f"Migrating {table_name}...")
|
||||
sqlite_cursor.execute(f"SELECT * FROM {table_name}")
|
||||
sqlite_cursor.execute(f"SELECT * FROM \"{table_name}\"")
|
||||
rows = sqlite_cursor.fetchall()
|
||||
|
||||
count = 0
|
||||
|
||||
@@ -11,9 +11,13 @@ from urllib.parse import urljoin
|
||||
|
||||
import aiohttp
|
||||
|
||||
from aitbc import get_logger
|
||||
|
||||
from .exceptions import *
|
||||
from .models import *
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class AgentIdentityClient:
|
||||
"""Main client for the AITBC Agent Identity SDK"""
|
||||
@@ -460,9 +464,9 @@ async def create_identity_with_wallets(
|
||||
failed_wallets = [w for w in wallet_results if not w.get("success", False)]
|
||||
|
||||
if failed_wallets:
|
||||
print(f"Warning: {len(failed_wallets)} wallets failed to create")
|
||||
logger.warning(f"{len(failed_wallets)} wallets failed to create")
|
||||
for wallet in failed_wallets:
|
||||
print(f" Chain {wallet['chain_id']}: {wallet.get('error', 'Unknown error')}")
|
||||
logger.warning(f"Chain {wallet['chain_id']}: {wallet.get('error', 'Unknown error')}")
|
||||
|
||||
return identity_response
|
||||
|
||||
@@ -505,7 +509,7 @@ async def verify_identity_on_all_chains(
|
||||
verification_results.append(result)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to verify on chain {mapping.chain_id}: {e}")
|
||||
logger.error(f"Failed to verify on chain {mapping.chain_id}: {e}")
|
||||
|
||||
return verification_results
|
||||
|
||||
|
||||
3
apps/coordinator-api/src/app/contexts/__init__.py
Normal file
3
apps/coordinator-api/src/app/contexts/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""Bounded contexts for the Coordinator API."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Agent Identity bounded context."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Agent Identity domain models."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -136,7 +136,7 @@ class CrossChainMapping(SQLModel, table=True):
|
||||
class IdentityVerification(SQLModel, table=True):
|
||||
"""Verification records for cross-chain identities"""
|
||||
|
||||
__tablename__ = "identity_verifications"
|
||||
__tablename__ = IDENTITY_VERIFICATION_TABLE
|
||||
__table_args__ = {"extend_existing": True}
|
||||
|
||||
id: str = Field(default_factory=lambda: f"verify_{uuid4().hex[:8]}", primary_key=True)
|
||||
@@ -0,0 +1,7 @@
|
||||
"""Agent Identity routers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .agent_identity import router as agent_identity
|
||||
|
||||
__all__ = ["agent_identity"]
|
||||
@@ -10,13 +10,13 @@ from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from fastapi.responses import JSONResponse
|
||||
from sqlmodel import Session
|
||||
|
||||
from ..agent_identity.manager import AgentIdentityManager
|
||||
from ..domain.agent_identity import (
|
||||
from ....agent_identity.manager import AgentIdentityManager
|
||||
from ....domain.agent_identity import (
|
||||
CrossChainMappingResponse,
|
||||
IdentityStatus,
|
||||
VerificationType,
|
||||
)
|
||||
from ..storage.db import get_session
|
||||
from ....storage.db import get_session
|
||||
|
||||
router = APIRouter(prefix="/agent-identity", tags=["Agent Identity"])
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Agent Identity services."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Agent Identity storage layer."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Agent Identity context database schema."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Table name prefixes for agent identity context
|
||||
AGENT_IDENTITY_TABLE_PREFIX = "agent_identity_"
|
||||
|
||||
# Agent Identity context table names
|
||||
AGENT_IDENTITY_TABLE = f"{AGENT_IDENTITY_TABLE_PREFIX}identity"
|
||||
IDENTITY_VERIFICATION_TABLE = f"{AGENT_IDENTITY_TABLE_PREFIX}verification"
|
||||
CROSS_CHAIN_MAPPING_TABLE = f"{AGENT_IDENTITY_TABLE_PREFIX}cross_chain_mapping"
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Blockchain bounded context."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Blockchain domain models."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,7 @@
|
||||
"""Blockchain routers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .blockchain import router as blockchain
|
||||
|
||||
__all__ = ["blockchain"]
|
||||
@@ -16,7 +16,7 @@ router = APIRouter(tags=["blockchain"])
|
||||
async def blockchain_status() -> dict[str, Any]:
|
||||
"""Get blockchain status."""
|
||||
try:
|
||||
from ..config import settings
|
||||
from ....config import settings
|
||||
|
||||
rpc_url = settings.blockchain_rpc_url.rstrip("/")
|
||||
client = AITBCHTTPClient(timeout=5.0)
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Blockchain services."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -8,7 +8,7 @@ from aitbc import get_logger, AITBCHTTPClient, NetworkError
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
from ..config import settings
|
||||
from ....config import settings
|
||||
|
||||
BLOCKCHAIN_RPC = "http://127.0.0.1:9080/rpc"
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Blockchain storage layer."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Blockchain context database schema."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Table name prefixes for blockchain context
|
||||
BLOCKCHAIN_TABLE_PREFIX = "blockchain_"
|
||||
|
||||
# Blockchain context table names
|
||||
BLOCKCHAIN_STATUS_TABLE = f"{BLOCKCHAIN_TABLE_PREFIX}status"
|
||||
BLOCKCHAIN_TRANSACTION_TABLE = f"{BLOCKCHAIN_TABLE_PREFIX}transaction"
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Marketplace bounded context."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Marketplace domain models."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -29,7 +29,7 @@ class MarketplaceOffer(SQLModel, table=True):
|
||||
|
||||
|
||||
class MarketplaceBid(SQLModel, table=True):
|
||||
__tablename__ = "marketplacebid"
|
||||
__tablename__ = MARKETPLACE_BID_TABLE
|
||||
__table_args__ = {"extend_existing": True}
|
||||
|
||||
id: str = Field(default_factory=lambda: uuid4().hex, primary_key=True)
|
||||
@@ -0,0 +1,17 @@
|
||||
"""Marketplace routers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .marketplace import router as marketplace
|
||||
from .marketplace_gpu import router as marketplace_gpu
|
||||
from .marketplace_offers import router as marketplace_offers
|
||||
from .global_marketplace import router as global_marketplace
|
||||
from .global_marketplace_integration import router as global_marketplace_integration
|
||||
|
||||
__all__ = [
|
||||
"marketplace",
|
||||
"marketplace_gpu",
|
||||
"marketplace_offers",
|
||||
"global_marketplace",
|
||||
"global_marketplace_integration",
|
||||
]
|
||||
@@ -9,8 +9,8 @@ from typing import Any
|
||||
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query
|
||||
from sqlmodel import Session, func, select
|
||||
|
||||
from ..agent_identity.manager import AgentIdentityManager
|
||||
from ..domain.global_marketplace import (
|
||||
from ....agent_identity.manager import AgentIdentityManager
|
||||
from ....domain.global_marketplace import (
|
||||
GlobalMarketplaceConfig,
|
||||
GlobalMarketplaceOffer,
|
||||
GlobalMarketplaceTransaction,
|
||||
@@ -18,8 +18,8 @@ from ..domain.global_marketplace import (
|
||||
MarketplaceStatus,
|
||||
RegionStatus,
|
||||
)
|
||||
from ..services.global_marketplace import GlobalMarketplaceService, RegionManager
|
||||
from ..storage.db import get_session
|
||||
from ....services.global_marketplace import GlobalMarketplaceService, RegionManager
|
||||
from ....storage.db import get_session
|
||||
|
||||
router = APIRouter(prefix="/global-marketplace", tags=["Global Marketplace"])
|
||||
|
||||
@@ -9,18 +9,18 @@ from typing import Any
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from ..agent_identity.manager import AgentIdentityManager
|
||||
from ..domain.global_marketplace import (
|
||||
from ....agent_identity.manager import AgentIdentityManager
|
||||
from ....domain.global_marketplace import (
|
||||
GlobalMarketplaceOffer,
|
||||
)
|
||||
from ..reputation.engine import CrossChainReputationEngine
|
||||
from ..services.cross_chain_bridge_enhanced import BridgeProtocol
|
||||
from ..services.global_marketplace_integration import (
|
||||
from ....reputation.engine import CrossChainReputationEngine
|
||||
from ....services.cross_chain_bridge_enhanced import BridgeProtocol
|
||||
from ....services.global_marketplace_integration import (
|
||||
GlobalMarketplaceIntegrationService,
|
||||
IntegrationStatus,
|
||||
)
|
||||
from ..services.multi_chain_transaction_manager import TransactionPriority
|
||||
from ..storage.db import get_session
|
||||
from ....services.multi_chain_transaction_manager import TransactionPriority
|
||||
from ....storage.db import get_session
|
||||
|
||||
router = APIRouter(prefix="/global-marketplace-integration", tags=["Global Marketplace Integration"])
|
||||
|
||||
@@ -6,12 +6,12 @@ from slowapi.util import get_remote_address
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from aitbc import get_logger
|
||||
from ..config import settings
|
||||
from ..metrics import marketplace_errors_total, marketplace_requests_total
|
||||
from ..schemas import MarketplaceBidRequest, MarketplaceBidView, MarketplaceOfferView, MarketplaceStatsView
|
||||
from ..services import MarketplaceService
|
||||
from ..storage import get_session
|
||||
from ..utils.cache import cached, get_cache_config
|
||||
from ....config import settings
|
||||
from ....metrics import marketplace_errors_total, marketplace_requests_total
|
||||
from ....schemas import MarketplaceBidRequest, MarketplaceBidView, MarketplaceOfferView, MarketplaceStatsView
|
||||
from ...services import MarketplaceService
|
||||
from ....storage import get_session
|
||||
from ....utils.cache import cached, get_cache_config
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -16,13 +16,13 @@ from sqlalchemy.orm import Session
|
||||
from sqlmodel import col, func, select
|
||||
|
||||
from aitbc import get_logger
|
||||
from ..custom_types import Constraints
|
||||
from ..domain.gpu_marketplace import GPUBooking, GPURegistry, GPUReview
|
||||
from ..domain.job import Job
|
||||
from ..schemas import JobCreate, JobPaymentCreate
|
||||
from ..services.dynamic_pricing_engine import DynamicPricingEngine, PricingStrategy, ResourceType
|
||||
from ..services.jobs import JobService
|
||||
from ..services.market_data_collector import MarketDataCollector
|
||||
from ....custom_types import Constraints
|
||||
from ....domain.gpu_marketplace import GPUBooking, GPURegistry, GPUReview
|
||||
from ....domain.job import Job
|
||||
from ....schemas import JobCreate, JobPaymentCreate
|
||||
from ....services.dynamic_pricing_engine import DynamicPricingEngine, PricingStrategy, ResourceType
|
||||
from ....services.jobs import JobService
|
||||
from ....services.market_data_collector import MarketDataCollector
|
||||
from ..services.payments import PaymentService
|
||||
from ..storage.db import get_session
|
||||
|
||||
@@ -12,10 +12,10 @@ from fastapi import APIRouter, Depends, HTTPException
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from aitbc import get_logger
|
||||
from ..deps import require_admin_key
|
||||
from ..domain import MarketplaceOffer, Miner
|
||||
from ..schemas import MarketplaceOfferView
|
||||
from ..storage import get_session
|
||||
from ....deps import require_admin_key
|
||||
from ....domain import MarketplaceOffer, Miner
|
||||
from ....schemas import MarketplaceOfferView
|
||||
from ....storage import get_session
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Marketplace services."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -4,8 +4,8 @@ from statistics import mean
|
||||
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from ..domain import MarketplaceBid, MarketplaceOffer
|
||||
from ..schemas import (
|
||||
from ....domain import MarketplaceBid, MarketplaceOffer
|
||||
from ....schemas import (
|
||||
MarketplaceBidRequest,
|
||||
MarketplaceBidView,
|
||||
MarketplaceOfferView,
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user