File size: 2,572 Bytes
f02fdcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import os
import asyncio
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from src.config.settings import Settings
from src.config.security import SecurityConfig
from src.database import init_db, MemoryModel, SessionLocal
from src.agents.router import RouterAgent
from src.agents.base import BaseAgent
from src.brain.dual import DualBrain
from src.utils.tokens import TokenManager
from src.utils.learner import Learner
app = FastAPI(
title="Super Brain v2",
version="2.0.0",
debug=Settings.DEBUG
)
# CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize database
init_db()
# Initialize components
router_agent = RouterAgent()
dual_brain = DualBrain()
token_manager = TokenManager()
@app.on_event("startup")
async def startup_event():
if not os.path.exists("data"):
os.makedirs("data")
print("Super Brain v2 started")
@app.on_event("shutdown")
async def shutdown_event():
print("Super Brain v2 shutting down")
@app.post("/api/chat")
async def chat(request: Request):
body = await request.json()
prompt = body.get("prompt", "")
# Route through security
req_dict = {
"url": request.url.path,
"ip": request.client.host,
"method": request.method,
"body": body
}
route_result = await router_agent.execute(req_dict)
if isinstance(route_result, tuple) and len(route_result) == 2:
result, status_code = route_result
if status_code != 200:
raise HTTPException(status_code=status_code, detail=result.get("error", "Error"))
# Process through brain
response = dual_brain.process(prompt)
# Store in memory
db = SessionLocal()
try:
memory_entry = MemoryModel(key=prompt, value=response)
db.add(memory_entry)
db.commit()
except Exception:
db.rollback()
finally:
db.close()
return {"response": response}
@app.post("/api/train")
async def train():
dataset_path = "brain_dataset.jsonl"
if os.path.exists(dataset_path):
Learner.train(dataset_path, dual_brain)
return {"status": "Training complete"}
return {"error": "Dataset not found"}, 404
@app.get("/api/health")
async def health():
return {"status": "ok", "version": "2.0.0"}
if __name__ == "__main__":
import uvicorn
uvicorn.run("src.server:app", host="0.0.0.0", port=8000, reload=Settings.DEBUG) |