diff --git a/backend/Dockerfile.python b/backend/Dockerfile.python new file mode 100644 index 00000000..3e850cf8 --- /dev/null +++ b/backend/Dockerfile.python @@ -0,0 +1,15 @@ +FROM python:3.12-slim +WORKDIR /app + +# Create non-root user +RUN adduser --disabled-password --gecos "" executor + +# Create execution directory with proper permissions +RUN mkdir /app/execution && chown executor:executor /app/execution + +# Switch to non-root user +USER executor +WORKDIR /app/execution + +# Keep container running +CMD ["tail", "-f", "/dev/null"] diff --git a/backend/openui/code_executor.py b/backend/openui/code_executor.py new file mode 100644 index 00000000..be66999c --- /dev/null +++ b/backend/openui/code_executor.py @@ -0,0 +1,102 @@ +from typing import Dict, Any +from pydantic import BaseModel +import docker +from docker.errors import APIError, ContainerError +from . import config + +class CodeExecutionRequest(BaseModel): + """Request model for code execution.""" + code: str + language: str = "python" # "python" or "bash" + timeout: int = 30 # seconds + +class CodeExecutionResponse(BaseModel): + """Response model for code execution results.""" + output: str + error: str | None = None + status: str + +class CodeExecutor: + """Handles code execution in isolated Docker containers.""" + + def __init__(self): + """Initialize Docker client.""" + try: + self.client = docker.from_env() + except Exception as e: + raise RuntimeError(f"Failed to initialize Docker client: {e}") + + async def execute(self, request: CodeExecutionRequest) -> CodeExecutionResponse: + """ + Execute code in an isolated Docker container. + + Args: + request: CodeExecutionRequest containing code and execution parameters + + Returns: + CodeExecutionResponse with execution results + """ + try: + # Configure container based on language + if request.language == "python": + command = ["python", "-c", request.code] + image = "python:3.12-slim" + elif request.language == "bash": + command = ["bash", "-c", request.code] + image = "ubuntu:22.04" + else: + return CodeExecutionResponse( + output="", + error=f"Unsupported language: {request.language}", + status="error" + ) + + # Run code in isolated container + container = self.client.containers.run( + image, + command, + remove=True, + detach=True, + network_disabled=True, + mem_limit="512m", + pids_limit=50, + stdout=True, + stderr=True + ) + + try: + # Wait for container to finish with timeout + container.wait(timeout=request.timeout) + # Get execution output + output = container.logs(stdout=True, stderr=True).decode() + return CodeExecutionResponse( + output=output, + status="success" + ) + finally: + try: + container.remove(force=True) + except: + pass # Container might already be removed + except ContainerError as e: + # Handle container execution errors + error_msg = str(e) + return CodeExecutionResponse( + output="", + error=error_msg, + status="error" + ) + except APIError as e: + # Handle Docker API errors + return CodeExecutionResponse( + output="", + error=str(e), + status="error" + ) + except Exception as e: + # Handle unexpected errors + return CodeExecutionResponse( + output="", + error=f"Unexpected error: {str(e)}", + status="error" + ) diff --git a/backend/openui/db/models.py b/backend/openui/db/models.py index 0e057842..bde77845 100644 --- a/backend/openui/db/models.py +++ b/backend/openui/db/models.py @@ -13,6 +13,7 @@ ) import uuid import datetime +from pathlib import Path from playhouse.sqlite_ext import SqliteExtDatabase, JSONField from playhouse.migrate import SqliteMigrator, migrate from openui import config @@ -34,6 +35,7 @@ class Meta: class SchemaMigration(BaseModel): + id = BinaryUUIDField(primary_key=True) version = CharField() @@ -81,6 +83,7 @@ class Usage(BaseModel): output_tokens = IntegerField() day = DateField() user = ForeignKeyField(User, backref="usage") + user_id = BinaryUUIDField() # Explicit field for foreign key class Meta: primary_key = CompositeKey("user", "day") @@ -122,13 +125,20 @@ def alter(schema: SchemaMigration, ops: list[list], version: str) -> bool: except OperationalError as e: print("Migration failed", e) return False - schema.version = version - schema.save() + # Update version through model API + schema.update(version=version).where(SchemaMigration.id == schema.id).execute() print(f"Migrated {version}") return version != CURRENT_VERSION def perform_migration(schema: SchemaMigration) -> bool: + """Perform database schema migration. + + Args: + schema: Current schema migration record + Returns: + bool: True if migration was performed, False otherwise + """ if schema.version == "2024-03-08": version = "2024-03-12" aaguid = CharField(null=True) @@ -142,18 +152,21 @@ def perform_migration(schema: SchemaMigration) -> bool: version, ) if altered: - perform_migration(schema) + return perform_migration(schema) + return True if schema.version == "2024-03-12": version = "2024-05-14" database.create_tables([Vote]) - schema.version = version - schema.save() + schema.update(version=version).where(SchemaMigration.id == schema.id).execute() if version != CURRENT_VERSION: - perform_migration(schema) + return perform_migration(schema) + return True + return False # No migration needed def ensure_migrated(): - if not config.DB.exists(): + db_path = Path(config.DB) + if not db_path.exists(): database.create_tables( [User, Credential, Session, Component, SchemaMigration, Usage, Vote] ) diff --git a/backend/openui/litellm.py b/backend/openui/litellm.py index 9486a01e..01d8d1b9 100644 --- a/backend/openui/litellm.py +++ b/backend/openui/litellm.py @@ -1,7 +1,7 @@ import yaml import os import tempfile -import openai +from openai._client import AsyncOpenAI as OpenAI from .logs import logger @@ -98,7 +98,7 @@ def generate_config(): ) if "OPENAI_COMPATIBLE_ENDPOINT" in os.environ: - client = openai.OpenAI( + client = OpenAI( api_key=os.getenv("OPENAI_COMPATIBLE_API_KEY"), base_url=os.getenv("OPENAI_COMPATIBLE_ENDPOINT"), ) diff --git a/backend/openui/openai.py b/backend/openui/openai.py index 27f71c4d..88ea49ad 100644 --- a/backend/openui/openai.py +++ b/backend/openui/openai.py @@ -1,8 +1,7 @@ import json -from openai import AsyncStream -from openai.types.chat import ( - ChatCompletionChunk, -) +from openai._streaming import AsyncStream +from openai.types.chat import ChatCompletionChunk +from typing import AsyncGenerator from .db.models import Usage @@ -11,7 +10,7 @@ async def openai_stream_generator( input_tokens: int, user_id: str, multiplier: int = 1, -): +) -> AsyncGenerator[str, None]: # async for chunk in subscription.response.aiter_bytes(): # yield chunk output_tokens = 0 diff --git a/backend/openui/server.py b/backend/openui/server.py index 0f321a94..0cad5891 100644 --- a/backend/openui/server.py +++ b/backend/openui/server.py @@ -39,13 +39,15 @@ from .util import storage from .util import get_git_user_email from . import config +from .code_executor import CodeExecutor, CodeExecutionRequest from pydantic import ValidationError from multiprocessing import Queue -from openai import AsyncOpenAI, APIStatusError, AsyncStream -from openai.types.chat import ( - ChatCompletionChunk, -) +from openai import AsyncOpenAI, APIError as APIStatusError +from openai.types.chat import ChatCompletion, ChatCompletionChunk +from openai.types.completion import Completion +from openai import AsyncStream from ollama import AsyncClient, ResponseError +from typing import Optional, Any, AsyncGenerator, Dict, List, Union from pathlib import Path from typing import Optional import traceback @@ -111,9 +113,14 @@ async def lifespan(app: FastAPI): ) async def chat_completions( request: Request, - # chat_request: CompletionCreateParams, # TODO: lots' fo weirdness here, just using raw json - # ctx: Any = Depends(weave_context), -): +) -> StreamingResponse: + """Handle chat completion requests with streaming responses. + + Args: + request: FastAPI request object + Returns: + StreamingResponse: Server-sent events stream + """ if request.session.get("user_id") is None: raise HTTPException(status_code=401, detail="Login required to use OpenUI") user_id = request.session["user_id"] @@ -134,11 +141,9 @@ async def chat_completions( if data.get("model").startswith("gpt"): if data["model"] == "gpt-4" or data["model"] == "gpt-4-32k": raise HTTPException(status=400, data="Model not supported") - response: AsyncStream[ - ChatCompletionChunk - ] = await openai.chat.completions.create( + response = await openai.chat.completions.create( **data, - ) + ) # Type inference handles AsyncStream[ChatCompletionChunk] # gpt-4 tokens are 20x more expensive multiplier = 20 if "gpt-4" in data["model"] else 1 return StreamingResponse( @@ -150,11 +155,9 @@ async def chat_completions( data["model"] = data["model"].replace("groq/", "") if groq is None: raise HTTPException(status=500, detail="Groq API key is not set.") - response: AsyncStream[ - ChatCompletionChunk - ] = await groq.chat.completions.create( + response = await groq.chat.completions.create( **data, - ) + ) # Type inference handles AsyncStream[ChatCompletionChunk] return StreamingResponse( openai_stream_generator(response, input_tokens, user_id, 1), media_type="text/event-stream", @@ -164,11 +167,9 @@ async def chat_completions( data["model"] = data["model"].replace("litellm/", "") if litellm is None: raise HTTPException(status=500, detail="LiteLLM API key is not set.") - response: AsyncStream[ - ChatCompletionChunk - ] = await litellm.chat.completions.create( + response = await litellm.chat.completions.create( **data, - ) + ) # Type inference handles AsyncStream[ChatCompletionChunk] return StreamingResponse( openai_stream_generator(response, input_tokens, user_id, 1), media_type="text/event-stream", @@ -191,11 +192,9 @@ async def chat_completions( ) gen = await ollama_stream_generator(response, data) else: - response: AsyncStream[ - ChatCompletionChunk - ] = await ollama_openai.chat.completions.create( + response = await ollama_openai.chat.completions.create( **data, - ) + ) # Type inference handles AsyncStream[ChatCompletionChunk] def gen(): return openai_stream_generator(response, input_tokens, user_id, 0) @@ -391,7 +390,8 @@ async def vote(request: Request, payload: VoteRequest): return payload -async def get_openai_models(): +async def get_openai_models() -> List[str]: + """Get list of supported OpenAI models.""" try: await openai.models.list() # We only support 3.5 and 4 for now @@ -401,34 +401,57 @@ async def get_openai_models(): return [] -async def get_ollama_models(): +async def get_ollama_models() -> List[Dict[str, Any]]: + """Get list of available Ollama models.""" try: - return (await ollama.list())["models"] + if ollama is None or not hasattr(ollama, 'list'): + return [] + response = await ollama.list() + if response is None or not isinstance(response, dict): + return [] + return response.get("models", []) except Exception: logger.warning("Couldn't connect to Ollama at %s", config.OLLAMA_HOST) return [] -async def get_groq_models(): +async def get_groq_models() -> List[Any]: + """Get list of available Groq models.""" try: + if groq is None or not hasattr(groq, 'models'): + return [] + models = await groq.models.list() + if models is None: + return [] return [ - d for d in (await groq.models.list()).data if not d.id.startswith("whisper") + d for d in models.data if not d.id.startswith("whisper") ] except Exception: logger.warning("Couldn't connect to Groq at %s", config.GROQ_BASE_URL) return [] -async def get_litellm_models(): +async def get_litellm_models() -> List[Any]: + """Get list of available LiteLLM models.""" try: - return (await litellm.models.list()).data + if litellm is None or not hasattr(litellm, 'models'): + return [] + models = await litellm.models.list() + if models is None: + return [] + return models.data except Exception: logger.warning("Couldn't connect to LiteLLM at %s", config.LITELLM_BASE_URL) return [] @router.get("/v1/models", tags="openui/models") -async def models(): +async def models() -> Dict[str, Dict[str, List[Union[str, Dict[str, Any]]]]]: + """Get available models from all providers. + + Returns: + Dict containing provider name to list of models mapping + """ tasks = [ get_openai_models(), get_groq_models(), @@ -492,6 +515,8 @@ async def get_session( else: raise HTTPException(status_code=404, detail="No session found") session_data = session_store.get(session_id) + if session_data is None: + raise HTTPException(status_code=404, detail="Session data not found") return JSONResponse( content=session_data.model_dump(), status_code=200, @@ -516,6 +541,28 @@ async def delete_session( ) +@router.post("/v1/execute", tags=["openui/execute"]) +async def execute_code(request: Request, code_request: CodeExecutionRequest): + """Execute code in an isolated container. + + Args: + request: FastAPI request object + code_request: Code execution parameters + + Returns: + CodeExecutionResponse with execution results + + Raises: + HTTPException: If user is not authenticated + """ + if request.session.get("user_id") is None: + raise HTTPException(status_code=401, detail="Login required") + + executor = CodeExecutor() + result = await executor.execute(code_request) + return result + + @router.get("/openui/{name}.svg", tags=["openui/svg"]) async def render_svg(name, text: Optional[str] = None): dims = name.split("x") diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 83d02ee6..cfa55575 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -4,11 +4,15 @@ dependencies = [ "openai>=1.12.0", "ollama>=0.1.7", "itsdangerous>=2.0.1", - "peewee>=3.0.0", + "peewee>=3.16.0", # Updated for better type support "fastapi>=0.111.0", "uvicorn>=0.22.0", "fastapi-sso>=0.10.0", "boto3>=1.34.67", + "types-peewee>=3.16.0", # Added type stubs + "docker>=7.0.0", # Added for code execution + "ollama>=0.1.6", # Added for Ollama integration + "tiktoken>=0.6.0", # Added for token counting ] name = "openui" version = "0.4.0" @@ -34,8 +38,12 @@ CI = "https://github.com/wandb/openui/actions" [project.optional-dependencies] -test = ["pytest", "pre-commit"] +test = ["pytest", "pre-commit", "mypy>=1.8.0"] # Added type checker litellm = ["litellm[proxy]>=1.40.20", "google-generativeai"] +types = [ # Added type stubs + "types-requests>=2.31.0", + "types-PyYAML>=6.0.12", +] eval = [ "beautifulsoup4>=4.0.0", "mistletoe>=1.0.0", diff --git a/docker-compose.yaml b/docker-compose.yaml index 0fe38693..c3b43f9d 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -20,8 +20,32 @@ services: - OPENAI_API_KEY - ANTHROPIC_API_KEY - GROQ_API_KEY + networks: + - code-execution env_file: - .env + volumes: + - /var/run/docker.sock:/var/run/docker.sock + group_add: + - ${DOCKER_GROUP_ID:-999} # Default Docker group ID on most systems + + code-executor: + build: + context: ./backend + dockerfile: Dockerfile.python + security_opt: + - no-new-privileges + read_only: true + tmpfs: + - /tmp + - /app/execution + networks: + - code-execution + restart: always + +networks: + code-execution: + name: code-execution volumes: ollama: diff --git a/frontend/src/components/CodeViewer.tsx b/frontend/src/components/CodeViewer.tsx index 4759385f..6c8adf51 100644 --- a/frontend/src/components/CodeViewer.tsx +++ b/frontend/src/components/CodeViewer.tsx @@ -17,13 +17,17 @@ import { convertFrameworkAtom, historyAtomFamily, selectedFrameworkAtom, - uiThemeAtom, - type Framework + uiThemeAtom } from 'state' + +import { Framework } from '../state/atoms/history' import { wrappedCode } from '../lib/html' import { cn } from '../lib/utils' import Scaffold from './Scaffold' + +const isExecutableLanguage = (f: Framework): boolean => ['python', 'javascript', 'bash'].includes(f) + const CodeEditor = lazy(async () => import('components/CodeEditor')) function replaceAll(str: string, mapObj: Record) { @@ -66,6 +70,13 @@ ${render interface ViewerProps { id: string code: string + framework?: string +} + +interface ExecutionResult { + output: string + error?: string + status: string } function stripCodeblocks(code: string) { @@ -78,6 +89,9 @@ export default function CodeViewer({ id, code }: ViewerProps) { const theme = themes.find(t => t.name === uiTheme) const [framework, setFramework] = useAtom(selectedFrameworkAtom) const [convertFramework, setConvertFramework] = useAtom(convertFrameworkAtom) + const handleFrameworkChange = (f: Framework) => { + setFramework(f) + } // local state const [currentCode, setCurrentCode] = useState(code) @@ -109,6 +123,33 @@ export default function CodeViewer({ id, code }: ViewerProps) { frameworks.push('jsx') } + const [executionResult, setExecutionResult] = useState(null) + const [isExecuting, setIsExecuting] = useState(false) + + const executeCode = async () => { + setIsExecuting(true) + try { + const response = await fetch('/v1/execute', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + code: currentCode, + language: isExecutableLanguage(framework) ? framework : 'javascript' + }) + }) + const result = await response.json() + setExecutionResult(result) + } catch (error) { + setExecutionResult({ + output: '', + error: String(error), + status: 'error' + }) + } finally { + setIsExecuting(false) + } + } + return (
@@ -119,7 +160,7 @@ export default function CodeViewer({ id, code }: ViewerProps) {
+
+ + {executionResult && ( +
+ 
+							{executionResult.error || executionResult.output}
+
+
+ )} +
) } diff --git a/frontend/src/components/HtmlAnnotator.tsx b/frontend/src/components/HtmlAnnotator.tsx index dc6f5766..97c43b3a 100644 --- a/frontend/src/components/HtmlAnnotator.tsx +++ b/frontend/src/components/HtmlAnnotator.tsx @@ -161,13 +161,13 @@ export default function HTMLAnnotator({ error, id }: HTMLAnnotatorProps) { const toggleFacet = useCallback( (name: string) => { - setFacets(prevFacets => { + setFacets((prevFacets: string[]) => { // Clear all non-official facets - const cleanFacets = prevFacets.filter(facet => + const cleanFacets = prevFacets.filter((facet: string) => adjectives.includes(facet) ) if (cleanFacets.includes(name)) { - return cleanFacets.filter(facet => facet !== name) + return cleanFacets.filter((facet: string) => facet !== name) } return [...cleanFacets, name] }) diff --git a/frontend/src/lib/markdown.ts b/frontend/src/lib/markdown.ts index 457ada8f..095d6881 100644 --- a/frontend/src/lib/markdown.ts +++ b/frontend/src/lib/markdown.ts @@ -80,20 +80,20 @@ export function parseMarkdown( } */ - const parsed = unified().use(remarkParse).parse(cleanMarkdown) + const parsed = unified().use(remarkParse).parse(cleanMarkdown) as { children: any[] } let htmlBlocks = parsed.children.filter( - c => - (c.type === 'code' && ['html', ''].includes(c.lang ?? '')) || + (c: any) => + (c.type === 'code' && ['html', ''].includes((c as Code).lang ?? '')) || c.type === 'html' ) as Code[] // TODO: maybe do this first and only if the first paragraph is chill - for (const c of parsed.children) { + for (const c of parsed.children as any[]) { if (c.type === 'paragraph') { let html = '' - if (c.children[0].type === 'html') { + if (c.children?.[0]?.type === 'html') { for (const c2 of c.children) { - html = html + (c2 as unknown as Code).value || '' + html = html + ((c2 as unknown as Code).value || '') } } htmlBlocks.push({ type: 'code', lang: 'html', value: html }) @@ -114,7 +114,7 @@ export function parseMarkdown( ] } const jsBlocks = parsed.children.filter( - c => c.type === 'code' && c.lang === 'javascript' + (c: any) => c.type === 'code' && (c as Code).lang === 'javascript' ) as Code[] result.html = fixHTML( [ diff --git a/frontend/src/state/atoms/history.ts b/frontend/src/state/atoms/history.ts index 925c4d88..fba658f7 100644 --- a/frontend/src/state/atoms/history.ts +++ b/frontend/src/state/atoms/history.ts @@ -13,6 +13,9 @@ export type Framework = | 'svelte' | 'vue' | 'web component' + | 'python' + | 'javascript' + | 'bash' export const FRAMEWORKS: Framework[] = [ 'preact', diff --git a/frontend/src/types/remark-parse.d.ts b/frontend/src/types/remark-parse.d.ts new file mode 100644 index 00000000..81a1dec4 --- /dev/null +++ b/frontend/src/types/remark-parse.d.ts @@ -0,0 +1,5 @@ +declare module 'remark-parse' { + import { Plugin } from 'unified' + const remarkParse: Plugin + export default remarkParse +} diff --git a/frontend/src/types/unified.d.ts b/frontend/src/types/unified.d.ts new file mode 100644 index 00000000..cef09efb --- /dev/null +++ b/frontend/src/types/unified.d.ts @@ -0,0 +1,12 @@ +declare module 'unified' { + export interface Plugin { + (...args: any[]): any + } + + export interface Processor { + use(plugin: Plugin, ...args: any[]): this + parse(content: string): any + } + + export function unified(): Processor +} diff --git a/frontend/tsconfig.json b/frontend/tsconfig.json index 9ae17efd..adda2a4b 100644 --- a/frontend/tsconfig.json +++ b/frontend/tsconfig.json @@ -14,9 +14,12 @@ "strict": true, "target": "ESNext", "paths": { - "@/*": ["./app/*"] + "@/*": ["./app/*"], + "components/*": ["./components/*"], + "lib/*": ["./lib/*"], + "state": ["./state"] }, - "types": ["vite/client", "vitest/globals", "vite-plugin-pwa/client"] + "types": ["vite/client", "vitest/globals", "vite-plugin-pwa/client", "react", "node"] }, "include": ["src"], "references": [{ "path": "./tsconfig.node.json" }]