Create main.py

main.py

@@ -0,0 +1,179 @@
+import os, json, asyncio, shutil, subprocess
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, BackgroundTasks
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse, StreamingResponse
+from pydantic import BaseModel
+from typing import List, Optional, Dict, Any
+from pathlib import Path
+import git
+
+from model_router import route_chat, route_autocomplete, get_available_models
+from agent import CodingAgent
+from codebase_index import CodebaseIndex
+from autocomplete import generate_autocomplete
+
+app = FastAPI(title="Codeki API")
+
+# Global codebase indexer (lazy load per project)
+indexers = {}
+
+# Agent instances for cloud tasks
+agents = {}
+
+# Mount static frontend
+if os.path.exists("static"):
+    app.mount("/static", StaticFiles(directory="static", html=True), name="static")
+
+# ---- Models ----
+class ChatRequest(BaseModel):
+    model: str = "groq-mixtral-8x7b-32768"
+    messages: List[Dict[str, str]]
+    context_files: Optional[List[str]] = None
+    project_path: str = "/app/sandbox"  # default
+
+class AgentTaskRequest(BaseModel):
+    prompt: str
+    model: str = "openrouter-anthropic/claude-3-opus"
+    project_path: str
+    auto_approve: bool = False
+    background: bool = False
+
+class AutocompleteRequest(BaseModel):
+    prefix: str
+    suffix: str = ""
+    filepath: str
+    language: str
+    max_tokens: int = 64
+    model: str = "groq-llama3-8b-8192"
+
+class IndexRequest(BaseModel):
+    project_path: str
+
+class BugFixRequest(BaseModel):
+    file_path: str
+    error_output: str
+    project_path: str
+
+class ReviewRequest(BaseModel):
+    pr_diff: str
+    project_path: str
+
+# ---- API Routes ----
+
+@app.get("/")
+async def root():
+    return FileResponse("static/index.html")
+
+@app.get("/api/models")
+async def get_models():
+    return get_available_models()
+
+@app.post("/api/chat")
+async def chat(req: ChatRequest):
+    # Context retrieval if files referenced
+    context = ""
+    if req.context_files:
+        for f in req.context_files:
+            p = Path(req.project_path) / f
+            if p.exists():
+                context += f"\n--- {f} ---\n{p.read_text()}\n"
+        req.messages.insert(0, {"role": "system", "content": f"Relevant codebase context:\n{context}"})
+    
+    async def event_stream():
+        async for token in route_chat(req.model, req.messages):
+            yield f"data: {json.dumps({'token': token})}\n\n"
+        yield "data: [DONE]\n\n"
+    return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+@app.post("/api/agent/run")
+async def run_agent(req: AgentTaskRequest, background_tasks: BackgroundTasks):
+    agent = CodingAgent(project_path=req.project_path, model=req.model)
+    if req.background:
+        # Cloud agent - run in background task queue
+        task_id = str(uuid.uuid4())
+        agents[task_id] = agent
+        background_tasks.add_task(agent.execute_task, req.prompt, auto_approve=req.auto_approve)
+        return {"task_id": task_id, "status": "started"}
+    else:
+        # Synchronous execution with streaming
+        async def event_stream():
+            async for event in agent.execute_task_stream(req.prompt, auto_approve=req.auto_approve):
+                yield f"data: {json.dumps(event)}\n\n"
+            yield "data: [DONE]\n\n"
+        return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+@app.post("/api/autocomplete")
+async def autocomplete(req: AutocompleteRequest):
+    suggestions = await generate_autocomplete(req)
+    return {"completions": suggestions}
+
+@app.post("/api/index")
+async def index_project(req: IndexRequest):
+    idx = CodebaseIndex(req.project_path)
+    idx.build_index()
+    indexers[req.project_path] = idx
+    return {"status": "indexed", "files": len(idx.file_list)}
+
+@app.post("/api/bugfix")
+async def bugfix(req: BugFixRequest):
+    # Agent that fixes the given error
+    agent = CodingAgent(project_path=req.project_path, model="openrouter-anthropic/claude-3-opus")
+    result = await agent.fix_error(req.file_path, req.error_output)
+    return result
+
+@app.post("/api/review")
+async def review_code(req: ReviewRequest):
+    agent = CodingAgent(project_path=req.project_path, model="openrouter-anthropic/claude-3-opus")
+    review = await agent.review_pr(req.pr_diff)
+    return {"review": review}
+
+@app.post("/api/git")
+async def git_operation(action: str, project_path: str, args: Optional[dict] = None):
+    try:
+        repo = git.Repo(project_path)
+        if action == "commit":
+            repo.git.add(A=True)
+            repo.git.commit(m=args.get("message", "auto commit"))
+        elif action == "log":
+            logs = list(repo.iter_commits(max_count=10))
+            return {"log": [{"message": c.message, "hash": c.hexsha} for c in logs]}
+        elif action == "branch":
+            branches = [b.name for b in repo.branches]
+            return {"branches": branches}
+        # ... other git actions
+    except Exception as e:
+        raise HTTPException(400, str(e))
+
+@app.websocket("/ws/terminal")
+async def terminal_ws(websocket: WebSocket):
+    await websocket.accept()
+    proc = None
+    try:
+        while True:
+            cmd = await websocket.receive_text()
+            if cmd.startswith("exec:"):
+                command = cmd[5:]
+                proc = await asyncio.create_subprocess_shell(
+                    command,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                    cwd="/app/sandbox"
+                )
+                stdout, stderr = await proc.communicate()
+                output = stdout.decode() + stderr.decode()
+                await websocket.send_text(output)
+            elif cmd == "ctrl-c" and proc:
+                proc.terminate()
+    except WebSocketDisconnect:
+        pass
+
+@app.post("/api/agent/task/{task_id}")
+async def get_agent_status(task_id: str):
+    agent = agents.get(task_id)
+    if not agent:
+        raise HTTPException(404, "Task not found")
+    return {"status": agent.status, "logs": agent.logs}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)