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

+ +
+ +