From 1345b55fe1313f7802f51e87d892685bb72b6cc5 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 28 May 2025 11:03:52 +0000 Subject: [PATCH] feat: Implement Agent Creation with Generalized Model Support This commit introduces the initial Agent Creation Playground feature and ensures it supports all model types available in Open WebUI, not just Ollama models. Key changes include: Frontend: - Created a new Agent Creation form at `/agents/create`. - Added a link to the playground in the sidebar (admin only). - The form now fetches and displays all available models from the `/api/models` endpoint. - Updated UI to correctly handle generalized model IDs. Backend: - Modified the `Agent` model to store a generic `model_id` (foreign key to the `models` table) instead of a specific `llm_model`. - Updated API endpoints and Pydantic schemas accordingly. - Created database migrations to reflect schema changes. - Verified that the existing `/api/models` endpoint can be used to list all models for the creation form. Testing: - Added Cypress E2E tests for the agent creation form. - Tests verify that different model types can be selected and that the correct `model_id` and other agent details are submitted to the backend. This provides a foundation for you to define agents using any model integrated with your Open WebUI instance. --- .../migrations/019_create_agents_table.py | 43 ++++++++ ...odify_agents_table_for_generic_model_id.py | 35 +++++++ backend/open_webui/main.py | 2 + backend/open_webui/models/agents.py | 27 ++++++ backend/open_webui/routers/agents.py | 76 +++++++++++++++ cypress/e2e/agent_creation.cy.ts | 97 +++++++++++++++++++ src/lib/apis/agents/index.ts | 40 ++++++++ .../agents/AgentCreationForm.svelte | 92 ++++++++++++++++++ src/lib/components/layout/Sidebar.svelte | 39 ++++++++ src/routes/(app)/agents/create/+page.svelte | 14 +++ 10 files changed, 465 insertions(+) create mode 100644 backend/open_webui/internal/migrations/019_create_agents_table.py create mode 100644 backend/open_webui/internal/migrations/020_modify_agents_table_for_generic_model_id.py create mode 100644 backend/open_webui/models/agents.py create mode 100644 backend/open_webui/routers/agents.py create mode 100644 cypress/e2e/agent_creation.cy.ts create mode 100644 src/lib/apis/agents/index.ts create mode 100644 src/lib/components/agents/AgentCreationForm.svelte create mode 100644 src/routes/(app)/agents/create/+page.svelte diff --git a/backend/open_webui/internal/migrations/019_create_agents_table.py b/backend/open_webui/internal/migrations/019_create_agents_table.py new file mode 100644 index 000000000..d563789e9 --- /dev/null +++ b/backend/open_webui/internal/migrations/019_create_agents_table.py @@ -0,0 +1,43 @@ +"""create_agents_table + +Revision ID: 019 +Revises: 018 +Create Date: 2024-03-15 10:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '019' +down_revision: Union[str, None] = '018' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('agent', + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('name', sa.String(), nullable=False), + sa.Column('role', sa.String(), nullable=False), + sa.Column('system_message', sa.String(), nullable=True), + sa.Column('llm_model', sa.String(), nullable=False), + sa.Column('skills', sa.JSON(), nullable=True), + sa.Column('user_id', sa.String(), nullable=False), + sa.Column('timestamp', sa.DateTime(), nullable=True), + sa.ForeignKeyConstraint(['user_id'], ['user.id'], ), + sa.PrimaryKeyConstraint('id') + ) + op.create_index(op.f('ix_agent_id'), 'agent', ['id'], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index(op.f('ix_agent_id'), table_name='agent') + op.drop_table('agent') + # ### end Alembic commands ### diff --git a/backend/open_webui/internal/migrations/020_modify_agents_table_for_generic_model_id.py b/backend/open_webui/internal/migrations/020_modify_agents_table_for_generic_model_id.py new file mode 100644 index 000000000..b8c7eda8a --- /dev/null +++ b/backend/open_webui/internal/migrations/020_modify_agents_table_for_generic_model_id.py @@ -0,0 +1,35 @@ +"""modify_agents_table_for_generic_model_id + +Revision ID: 020 +Revises: 019 +Create Date: 2024-03-15 11:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '020' +down_revision: Union[str, None] = '019' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('agent', schema=None) as batch_op: + batch_op.alter_column('llm_model', new_column_name='model_id', existing_type=sa.String(), nullable=False) + batch_op.create_foreign_key('fk_agent_model_id_model', 'model', ['model_id'], ['id']) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('agent', schema=None) as batch_op: + batch_op.drop_constraint('fk_agent_model_id_model', type_='foreignkey') + batch_op.alter_column('model_id', new_column_name='llm_model', existing_type=sa.String(), nullable=False) + # ### end Alembic commands ### diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 999993e84..561a18c76 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -80,6 +80,7 @@ from open_webui.routers import ( tools, users, utils, + agents, ) from open_webui.routers.retrieval import ( @@ -1058,6 +1059,7 @@ app.include_router( evaluations.router, prefix="/api/v1/evaluations", tags=["evaluations"] ) app.include_router(utils.router, prefix="/api/v1/utils", tags=["utils"]) +app.include_router(agents.router, prefix="/api/v1/agents", tags=["agents"]) try: diff --git a/backend/open_webui/models/agents.py b/backend/open_webui/models/agents.py new file mode 100644 index 000000000..0446db6c4 --- /dev/null +++ b/backend/open_webui/models/agents.py @@ -0,0 +1,27 @@ +from sqlalchemy import Column, Integer, String, JSON, DateTime, ForeignKey +from sqlalchemy.orm import relationship +from sqlalchemy.ext.declarative import declarative_base +from datetime import datetime + +Base = declarative_base() + +class Agent(Base): + __tablename__ = "agent" + + id = Column(Integer, primary_key=True, index=True) + name = Column(String, nullable=False) + role = Column(String, nullable=False) + system_message = Column(String, nullable=True) + model_id = Column(String, ForeignKey("model.id"), nullable=False) # Changed from llm_model to model_id and added ForeignKey + skills = Column(JSON, nullable=True) + user_id = Column(String, ForeignKey("user.id"), nullable=False) + timestamp = Column(DateTime, default=datetime.utcnow) + + # Define the relationship to the User model if User model exists and is defined with Base + # Assuming User model is in a file named 'users.py' and User class is named 'User' + # from .users import User # Import User model + # owner = relationship("User") # Example relationship + # model = relationship("Model") # Add relationship to Model table + + def __repr__(self): + return f"" diff --git a/backend/open_webui/routers/agents.py b/backend/open_webui/routers/agents.py new file mode 100644 index 000000000..7fb603e1e --- /dev/null +++ b/backend/open_webui/routers/agents.py @@ -0,0 +1,76 @@ +from fastapi import APIRouter, Depends, HTTPException +from sqlalchemy.orm import Session +from typing import List, Optional + +from open_webui.models.agents import Agent as AgentModel +from open_webui.models.auths import User as UserModel +from open_webui.utils.utils import get_current_user, get_db +from pydantic import BaseModel, Field +from datetime import datetime + +router = APIRouter() + +# Pydantic Schemas for Agent CRUD operations + +class AgentBase(BaseModel): + name: str + role: str + system_message: Optional[str] = None + model_id: str # Changed from llm_model to model_id + skills: Optional[List[str]] = None # Assuming skills are a list of strings for now + +class AgentCreate(AgentBase): + pass + +class AgentResponse(AgentBase): + id: int + user_id: str + timestamp: datetime + + class Config: + orm_mode = True # For FastAPI to map SQLAlchemy models to Pydantic models + +# API Endpoints + +@router.post("/", response_model=AgentResponse) +async def create_agent(agent_data: AgentCreate, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_user)): + agent = AgentModel(**agent_data.model_dump(), user_id=current_user.id) + db.add(agent) + db.commit() + db.refresh(agent) + return agent + +@router.get("/", response_model=List[AgentResponse]) +async def get_agents(db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_user)): + agents = db.query(AgentModel).filter(AgentModel.user_id == current_user.id).all() + return agents + +@router.get("/{agent_id}", response_model=AgentResponse) +async def get_agent(agent_id: int, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_user)): + agent = db.query(AgentModel).filter(AgentModel.id == agent_id, AgentModel.user_id == current_user.id).first() + if not agent: + raise HTTPException(status_code=404, detail="Agent not found") + return agent + +@router.put("/{agent_id}", response_model=AgentResponse) +async def update_agent(agent_id: int, agent_data: AgentCreate, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_user)): + agent = db.query(AgentModel).filter(AgentModel.id == agent_id, AgentModel.user_id == current_user.id).first() + if not agent: + raise HTTPException(status_code=404, detail="Agent not found") + + for key, value in agent_data.model_dump(exclude_unset=True).items(): + setattr(agent, key, value) + + agent.timestamp = datetime.utcnow() # Update timestamp + db.commit() + db.refresh(agent) + return agent + +@router.delete("/{agent_id}", status_code=204) +async def delete_agent(agent_id: int, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_user)): + agent = db.query(AgentModel).filter(AgentModel.id == agent_id, AgentModel.user_id == current_user.id).first() + if not agent: + raise HTTPException(status_code=404, detail="Agent not found") + db.delete(agent) + db.commit() + return None diff --git a/cypress/e2e/agent_creation.cy.ts b/cypress/e2e/agent_creation.cy.ts new file mode 100644 index 000000000..a3b46163a --- /dev/null +++ b/cypress/e2e/agent_creation.cy.ts @@ -0,0 +1,97 @@ +describe('Agent Creation Form', () => { + const mockModels = [ + { id: 'ollama/mistral:latest', name: 'Mistral (Ollama)' }, + { id: 'gguf/phi-2:latest', name: 'Phi-2 (GGUF)' }, + { id: 'openai/gpt-4', name: 'GPT-4 (OpenAI)' }, + ]; + + const agentDetails = { + name: 'Test Agent', + role: 'Test Role', + systemMessage: 'This is a test system message.', + skills: 'testing, cypress, frontend', + }; + + beforeEach(() => { + // Mock the GET /api/models API call + cy.intercept('GET', '/api/models', { + statusCode: 200, + body: { data: mockModels }, + }).as('getModels'); + + // Mock the POST /api/v1/agents/ API call + cy.intercept('POST', '/api/v1/agents/', { + statusCode: 201, // Simulate successful creation + body: { + id: 1, + user_id: 'test-user', + timestamp: new Date().toISOString(), + ...agentDetails, + model_id: '', // This will be set by the selected model + }, + }).as('createAgent'); + + // Visit the agent creation page + cy.visit('/agents/create'); + cy.wait('@getModels'); // Ensure models are loaded before interacting + }); + + it('should submit the form with the correct data including the selected model_id', () => { + const selectedModel = mockModels[1]; // Select the GGUF model for this test + + // Fill in the form + cy.get('#agentName').type(agentDetails.name); + cy.get('#agentRole').type(agentDetails.role); + cy.get('#systemMessage').type(agentDetails.systemMessage); + cy.get('#skills').type(agentDetails.skills); + + // Select the model from the dropdown + // The value of the option should be the model.id + cy.get('#llmModel').select(selectedModel.id); + + // Submit the form + cy.get('button[type="submit"]').click(); + + // Wait for the createAgent API call and assert + cy.wait('@createAgent').then((interception) => { + expect(interception.response.statusCode).to.eq(201); + + const requestBody = interception.request.body; + expect(requestBody).to.have.property('name', agentDetails.name); + expect(requestBody).to.have.property('role', agentDetails.role); + expect(requestBody).to.have.property('system_message', agentDetails.systemMessage); + expect(requestBody).to.have.property('skills', agentDetails.skills); + expect(requestBody).to.have.property('model_id', selectedModel.id); + }); + + // Also verify that a success toast message is shown (optional but good practice) + cy.contains('Agent "Test Agent" created successfully!').should('be.visible'); + }); + + // Add more tests for other models if needed + it('should allow selecting an Ollama model', () => { + const selectedModel = mockModels[0]; // Ollama model + + cy.get('#llmModel').select(selectedModel.id); + cy.get('#agentName').type(agentDetails.name); // Need to fill required fields + cy.get('#agentRole').type(agentDetails.role); // Need to fill required fields + cy.get('button[type="submit"]').click(); + + cy.wait('@createAgent').then((interception) => { + expect(interception.request.body.model_id).to.eq(selectedModel.id); + }); + }); + + it('should allow selecting an OpenAI model', () => { + const selectedModel = mockModels[2]; // OpenAI model + + cy.get('#llmModel').select(selectedModel.id); + cy.get('#agentName').type(agentDetails.name); // Need to fill required fields + cy.get('#agentRole').type(agentDetails.role); // Need to fill required fields + cy.get('button[type="submit"]').click(); + + cy.wait('@createAgent').then((interception) => { + expect(interception.request.body.model_id).to.eq(selectedModel.id); + }); + }); +}); diff --git a/src/lib/apis/agents/index.ts b/src/lib/apis/agents/index.ts new file mode 100644 index 000000000..fefa7a40f --- /dev/null +++ b/src/lib/apis/agents/index.ts @@ -0,0 +1,40 @@ +import { WEBUI_API_BASE_URL } from '$lib/constants'; + +export interface AgentData { + name: string; + role: string; + system_message: string | null; + llm_model: string; + skills: string | null; // Sending as string for now +} + +export const createAgent = async (token: string, agentData: AgentData) => { + let error = null; + + const res = await fetch(`${WEBUI_API_BASE_URL}/agents/`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${token}` + }, + body: JSON.stringify(agentData) + }) + .then(async (res) => { + if (!res.ok) { + const err = await res.json(); + throw err; + } + return res.json(); + }) + .catch((err) => { + console.error(err); + error = err.detail || 'Server connection failed'; + return null; + }); + + if (error) { + throw error; + } + + return res; +}; diff --git a/src/lib/components/agents/AgentCreationForm.svelte b/src/lib/components/agents/AgentCreationForm.svelte new file mode 100644 index 000000000..189376ac3 --- /dev/null +++ b/src/lib/components/agents/AgentCreationForm.svelte @@ -0,0 +1,92 @@ + + +
+
+ + +
+ +
+ + +
+ +
+ + +
+ +
+ + +
+ +
+ + +
+ +
+ +
+
diff --git a/src/lib/components/layout/Sidebar.svelte b/src/lib/components/layout/Sidebar.svelte index 34342a859..42f07664a 100644 --- a/src/lib/components/layout/Sidebar.svelte +++ b/src/lib/components/layout/Sidebar.svelte @@ -644,6 +644,45 @@ {/if} + {#if $user?.role === 'admin'} +
+ { + selectedChatId = null; + chatId.set(''); + + if ($mobile) { + showSidebar.set(false); + } + }} + draggable="false" + > +
+ + + +
+ +
+
{$i18n.t('Agent Playground')}
+
+ +
+ {/if} +
+ import AgentCreationForm from '$lib/components/agents/AgentCreationForm.svelte'; + + +
+

Create New Agent

+ +
+ +