""" AI API routes for code assistance """ import time from typing import List from fastapi import APIRouter, HTTPException, BackgroundTasks from fastapi.responses import StreamingResponse from app.models.ai import ( AIRequest, AIResponse, CodeCompletionRequest, CodeCompletionResponse, ExplainRequest, RefactorRequest, FixRequest, MultiFileRequest ) from app.services.ollama_service import ollama_service router = APIRouter() @router.post("/process", response_model=AIResponse) async def process_ai_request(request: AIRequest): """Process general AI request for code assistance""" start_time = time.time() try: # Check if Ollama is available if not await ollama_service.is_available(): raise HTTPException( status_code=503, detail="Ollama service is not available. Please ensure Ollama is running." ) # Build prompt based on command prompt = ollama_service.build_prompt( command=request.command, code=request.code, language=request.language, context=request.context ) # Generate completion result = await ollama_service.generate_completion( prompt=prompt, model=request.model ) execution_time = time.time() - start_time return AIResponse( success=True, result=result, execution_time=execution_time, model_used=request.model or ollama_service.default_model ) except Exception as e: execution_time = time.time() - start_time return AIResponse( success=False, result=f"Error processing request: {str(e)}", execution_time=execution_time ) @router.post("/explain", response_model=AIResponse) async def explain_code(request: ExplainRequest): """Explain code functionality""" start_time = time.time() try: if not await ollama_service.is_available(): raise HTTPException(status_code=503, detail="Ollama service unavailable") # Build explanation prompt prompt = f""" Explain the following {request.language.value if request.language else 'code'} code in {request.detail_level} detail: ```{request.language.value if request.language else 'code'} {request.code} ``` Please provide a clear explanation that covers: 1. What this code does 2. How it works 3. Key concepts used 4. Any potential issues or improvements Explanation:""" result = await ollama_service.generate_completion(prompt=prompt) execution_time = time.time() - start_time return AIResponse( success=True, result=result, execution_time=execution_time, model_used=ollama_service.default_model ) except Exception as e: execution_time = time.time() - start_time return AIResponse( success=False, result=f"Error explaining code: {str(e)}", execution_time=execution_time ) @router.post("/refactor", response_model=AIResponse) async def refactor_code(request: RefactorRequest): """Refactor code for better quality""" start_time = time.time() try: if not await ollama_service.is_available(): raise HTTPException(status_code=503, detail="Ollama service unavailable") prompt = f""" Refactor the following {request.language.value if request.language else 'code'} code to improve readability, performance, and maintainability: ```{request.language.value if request.language else 'code'} {request.code} ``` Focus on {request.refactor_type} improvements. Please provide: 1. The refactored code 2. Explanation of changes made 3. Benefits of the refactoring Refactored code:""" result = await ollama_service.generate_completion(prompt=prompt) execution_time = time.time() - start_time return AIResponse( success=True, result=result, execution_time=execution_time, model_used=ollama_service.default_model ) except Exception as e: execution_time = time.time() - start_time return AIResponse( success=False, result=f"Error refactoring code: {str(e)}", execution_time=execution_time ) @router.post("/fix", response_model=AIResponse) async def fix_code(request: FixRequest): """Fix bugs in code""" start_time = time.time() try: if not await ollama_service.is_available(): raise HTTPException(status_code=503, detail="Ollama service unavailable") prompt = ollama_service.build_prompt( command="fix", code=request.code, language=request.language, error_message=request.error_message ) result = await ollama_service.generate_completion(prompt=prompt) execution_time = time.time() - start_time return AIResponse( success=True, result=result, execution_time=execution_time, model_used=ollama_service.default_model ) except Exception as e: execution_time = time.time() - start_time return AIResponse( success=False, result=f"Error fixing code: {str(e)}", execution_time=execution_time ) @router.post("/complete", response_model=CodeCompletionResponse) async def complete_code(request: CodeCompletionRequest): """Generate code completion""" start_time = time.time() try: if not await ollama_service.is_available(): raise HTTPException(status_code=503, detail="Ollama service unavailable") # Extract context around cursor position code_before = request.code[:request.cursor_position] code_after = request.code[request.cursor_position:] prompt = f""" Complete the following {request.language.value if request.language else 'code'} code at the cursor position: ```{request.language.value if request.language else 'code'} {code_before}{code_after} ``` Provide only the code that should be inserted at the cursor position. Keep it concise and contextually appropriate. Completion:""" result = await ollama_service.generate_completion( prompt=prompt, max_tokens=request.max_tokens ) execution_time = time.time() - start_time return CodeCompletionResponse( success=True, completions=[result.strip()], cursor_position=request.cursor_position, execution_time=execution_time ) except Exception as e: execution_time = time.time() - start_time return CodeCompletionResponse( success=False, completions=[], cursor_position=request.cursor_position, execution_time=execution_time ) @router.post("/stream") async def stream_ai_response(request: AIRequest): """Stream AI response for real-time feedback""" try: if not await ollama_service.is_available(): raise HTTPException(status_code=503, detail="Ollama service unavailable") prompt = ollama_service.build_prompt( command=request.command, code=request.code, language=request.language, context=request.context ) async def generate(): try: async for chunk in ollama_service.generate_streaming( prompt=prompt, model=request.model ): yield f"data: {chunk}\n\n" yield "data: [DONE]\n\n" except Exception as e: yield f"data: ERROR: {str(e)}\n\n" return StreamingResponse( generate(), media_type="text/plain", headers={"Cache-Control": "no-cache"} ) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @router.post("/multifile", response_model=AIResponse) async def process_multifile_request(request: MultiFileRequest): """Process multi-file analysis request""" start_time = time.time() try: if not await ollama_service.is_available(): raise HTTPException(status_code=503, detail="Ollama service unavailable") # Build context from multiple files file_context = "\n\n".join([ f"File: {file_info['path']}\n```\n{file_info['content']}\n```" for file_info in request.files ]) # Focus on specific file if provided focus_context = "" if request.focus_file: focus_file = next( (f for f in request.files if f['path'] == request.focus_file), None ) if focus_file: focus_context = f"\n\nFocus on this file: {focus_file['path']}\n```\n{focus_file['content']}\n```" prompt = f""" Analyze the following multi-file codebase and {request.command.value}: {file_context} {focus_context} {f"Additional context: {request.context}" if request.context else ""} Please provide analysis considering the relationships between files and overall code structure. Response:""" result = await ollama_service.generate_completion(prompt=prompt) execution_time = time.time() - start_time return AIResponse( success=True, result=result, execution_time=execution_time, model_used=ollama_service.default_model, metadata={"files_analyzed": len(request.files)} ) except Exception as e: execution_time = time.time() - start_time return AIResponse( success=False, result=f"Error processing multifile request: {str(e)}", execution_time=execution_time ) @router.get("/health") async def ai_health(): """Check AI service health""" is_available = await ollama_service.is_available() models = await ollama_service.list_models() if is_available else [] return { "ollama_available": is_available, "models_count": len(models), "default_model": ollama_service.default_model, "base_url": ollama_service.base_url }