File size: 4,684 Bytes
bd91486 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Depends
from fastapi.security import OAuth2PasswordBearer
from pydantic import BaseModel
from typing import Optional, Dict, Any
from ..core.model import CogenBAI
from ..languages.generator import LanguageGenerator
from ..collaboration.session import SessionManager
from ..collaboration.websocket import collaboration_manager
from ..review.analyzer import CodeReviewAnalyzer
from ..testing.generator import TestGenerator
app = FastAPI(title="COGENBAI API")
model = CogenBAI()
lang_generator = LanguageGenerator()
session_manager = SessionManager()
code_reviewer = CodeReviewAnalyzer()
test_generator = TestGenerator()
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
class CodeRequest(BaseModel):
prompt: str
language: str
framework: Optional[str] = None
max_length: Optional[int] = 1024
temperature: float = 0.7
@app.post("/generate")
async def generate_code(request: CodeRequest, token: str = Depends(oauth2_scheme)) -> Dict[str, Any]:
try:
code = model.generate_code(
prompt=request.prompt,
language=request.language,
max_length=request.max_length,
temperature=request.temperature
)
return {"status": "success", "code": code}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/supported-languages")
async def get_supported_languages():
return {"languages": list(lang_generator.language_configs.keys())}
@app.websocket("/ws/{session_id}/{user_id}")
async def websocket_endpoint(websocket: WebSocket, session_id: str, user_id: str):
await collaboration_manager.connect(session_id, websocket)
try:
session = await session_manager.join_session(session_id, user_id)
if not session:
await websocket.close()
return
while True:
data = await websocket.receive_json()
if data["type"] == "code_update":
await session_manager.update_code(session_id, data["code"])
await collaboration_manager.broadcast(session_id, {
"type": "code_update",
"code": data["code"],
"user_id": user_id
})
except WebSocketDisconnect:
await collaboration_manager.disconnect(session_id, websocket)
@app.post("/sessions/create")
async def create_session(user_id: str):
session_id = f"session_{len(session_manager.sessions) + 1}"
session = await session_manager.create_session(session_id, user_id)
return session.to_dict()
@app.post("/review")
async def review_code(code: str, language: str) -> Dict[str, Any]:
try:
review_results = code_reviewer.review_code(code, language)
return {"status": "success", "review": review_results}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/generate-tests")
async def generate_tests(code: str, language: str, test_type: str = 'unit') -> Dict[str, str]:
try:
tests = test_generator.generate_tests(code, language, test_type)
return {"status": "success", "tests": tests}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/projects/create")
async def create_project(
name: str,
language: str,
framework: str,
initial_description: str
) -> Dict[str, Any]:
project_id = f"proj_{int(time.time())}"
project = ProjectState(
project_id=project_id,
name=name,
language=language,
framework=framework,
status="active",
completion_percentage=0.0,
last_modified=datetime.now(),
code_snippets={},
dependencies=[]
)
if model.project_tracker.create_project(project):
initial_code = model.generate_code(initial_description, language, framework)
project.code_snippets["initial"] = initial_code
model.project_tracker.update_project(
project_id,
{"code_snippets": json.dumps(project.code_snippets)}
)
return {"status": "success", "project_id": project_id, "code": initial_code}
raise HTTPException(status_code=500, detail="Failed to create project")
@app.post("/projects/{project_id}/continue")
async def continue_project(
project_id: str,
feature_description: str
) -> Dict[str, Any]:
try:
new_code = model.continue_project(project_id, feature_description)
return {"status": "success", "code": new_code}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
|