Add a websocket for streaming from the chat UI

pyproject.toml

@@ -76,6 +76,7 @@ dependencies = [
     "django-phonenumber-field == 7.3.0",
     "phonenumbers == 8.13.27",
     "markdownify ~= 0.11.6",
+    "websockets == 12.0",
 ]
 dynamic = ["version"]
 

src/khoj/configure.py

@@ -21,6 +21,7 @@
 from starlette.requests import HTTPConnection
 
 from khoj.database.adapters import (
+    AgentAdapters,
     ClientApplicationAdapters,
     ConversationAdapters,
     SubscriptionState,
@@ -229,11 +230,16 @@ def configure_server(
 
         state.SearchType = configure_search_types()
         state.search_models = configure_search(state.search_models, state.config.search_type)
+        setup_default_agent()
         initialize_content(regenerate, search_type, init, user)
     except Exception as e:
         raise e
 
 
+def setup_default_agent():
+    AgentAdapters.create_default_agent()
+
+
 def initialize_content(regenerate: bool, search_type: Optional[SearchType] = None, init=False, user: KhojUser = None):
     # Initialize Content from Config
     if state.search_models:
@@ -262,13 +268,15 @@ def initialize_content(regenerate: bool, search_type: Optional[SearchType] = Non
 def configure_routes(app):
     # Import APIs here to setup search types before while configuring server
     from khoj.routers.api import api
+    from khoj.routers.api_agents import api_agents
     from khoj.routers.api_chat import api_chat
     from khoj.routers.api_config import api_config
     from khoj.routers.indexer import indexer
     from khoj.routers.web_client import web_client
 
     app.include_router(api, prefix="/api")
     app.include_router(api_chat, prefix="/api/chat")
+    app.include_router(api_agents, prefix="/api/agents")
     app.include_router(api_config, prefix="/api/config")
     app.include_router(indexer, prefix="/api/v1/index")
     app.include_router(web_client)

src/khoj/database/adapters/__init__.py

@@ -16,6 +16,7 @@
 from torch import Tensor
 
 from khoj.database.models import (
+    Agent,
     ChatModelOptions,
     ClientApplication,
     Conversation,
@@ -37,6 +38,7 @@
     UserRequests,
     UserSearchModelConfig,
 )
+from khoj.processor.conversation import prompts
 from khoj.search_filter.date_filter import DateFilter
 from khoj.search_filter.file_filter import FileFilter
 from khoj.search_filter.word_filter import WordFilter
@@ -391,6 +393,79 @@ async def aget_client_application_by_id(client_id: str, client_secret: str):
         return await ClientApplication.objects.filter(client_id=client_id, client_secret=client_secret).afirst()
 
 
+class AgentAdapters:
+    DEFAULT_AGENT_NAME = "khoj"
+    DEFAULT_AGENT_AVATAR = "https://khoj-web-bucket.s3.amazonaws.com/lamp-128.png"
+
+    @staticmethod
+    async def aget_agent_by_id(agent_id: int):
+        return await Agent.objects.filter(id=agent_id).afirst()
+
+    @staticmethod
+    async def aget_agent_by_slug(agent_slug: str):
+        return await Agent.objects.filter(slug__iexact=agent_slug.lower()).afirst()
+
+    @staticmethod
+    def get_agent_by_slug(slug: str, user: KhojUser = None):
+        agent = Agent.objects.filter(slug=slug).first()
+        # Check if agent is public or created by the user
+        if agent and (agent.public or agent.creator == user):
+            return agent
+        return None
+
+    @staticmethod
+    def get_all_accessible_agents(user: KhojUser = None):
+        return Agent.objects.filter(Q(public=True) | Q(creator=user)).distinct().order_by("created_at")
+
+    @staticmethod
+    async def aget_all_accessible_agents(user: KhojUser = None) -> List[Agent]:
+        get_all_accessible_agents = sync_to_async(
+            lambda: Agent.objects.filter(Q(public=True) | Q(creator=user)).distinct().order_by("created_at").all(),
+            thread_sensitive=True,
+        )
+        agents = await get_all_accessible_agents()
+        return await sync_to_async(list)(agents)
+
+    @staticmethod
+    def get_conversation_agent_by_id(agent_id: int):
+        agent = Agent.objects.filter(id=agent_id).first()
+        if agent == AgentAdapters.get_default_agent():
+            # If the agent is set to the default agent, then return None and let the default application code be used
+            return None
+        return agent
+
+    @staticmethod
+    def get_default_agent():
+        return Agent.objects.filter(name=AgentAdapters.DEFAULT_AGENT_NAME).first()
+
+    @staticmethod
+    def create_default_agent():
+        default_conversation_config = ConversationAdapters.get_default_conversation_config()
+        default_personality = prompts.personality.format(current_date="placeholder")
+
+        if Agent.objects.filter(name=AgentAdapters.DEFAULT_AGENT_NAME).exists():
+            agent = Agent.objects.filter(name=AgentAdapters.DEFAULT_AGENT_NAME).first()
+            agent.tuning = default_personality
+            agent.chat_model = default_conversation_config
+            agent.save()
+            return agent
+
+        # The default agent is public and managed by the admin. It's handled a little differently than other agents.
+        return Agent.objects.create(
+            name=AgentAdapters.DEFAULT_AGENT_NAME,
+            public=True,
+            managed_by_admin=True,
+            chat_model=default_conversation_config,
+            tuning=default_personality,
+            tools=["*"],
+            avatar=AgentAdapters.DEFAULT_AGENT_AVATAR,
+        )
+
+    @staticmethod
+    async def aget_default_agent():
+        return await Agent.objects.filter(name=AgentAdapters.DEFAULT_AGENT_NAME).afirst()
+
+
 class ConversationAdapters:
     @staticmethod
     def get_conversation_by_user(
@@ -431,7 +506,14 @@ def get_conversation_by_id(conversation_id: int):
         return Conversation.objects.filter(id=conversation_id).first()
 
     @staticmethod
-    async def acreate_conversation_session(user: KhojUser, client_application: ClientApplication = None):
+    async def acreate_conversation_session(
+        user: KhojUser, client_application: ClientApplication = None, agent_slug: str = None
+    ):
+        if agent_slug:
+            agent = await AgentAdapters.aget_agent_by_slug(agent_slug)
+            if agent is None:
+                raise HTTPException(status_code=400, detail="Invalid agent id")
+            return await Conversation.objects.acreate(user=user, client=client_application, agent=agent)
         return await Conversation.objects.acreate(user=user, client=client_application)
 
     @staticmethod
@@ -446,7 +528,7 @@ async def aget_conversation_by_user(
             conversation = Conversation.objects.filter(user=user, client=client_application).order_by("-updated_at")
 
         if await conversation.aexists():
-            return await conversation.afirst()
+            return await conversation.prefetch_related("agent").afirst()
 
         return await Conversation.objects.acreate(user=user, client=client_application, slug=slug)
 
@@ -606,9 +688,14 @@ async def aget_conversation_starters(user: KhojUser, max_results=3):
         return random.sample(all_questions, max_results)
 
     @staticmethod
-    def get_valid_conversation_config(user: KhojUser):
+    def get_valid_conversation_config(user: KhojUser, conversation: Conversation):
         offline_chat_config = ConversationAdapters.get_offline_chat_conversation_config()
-        conversation_config = ConversationAdapters.get_conversation_config(user)
+
+        if conversation.agent and conversation.agent.chat_model:
+            conversation_config = conversation.agent.chat_model
+        else:
+            conversation_config = ConversationAdapters.get_conversation_config(user)
+
         if conversation_config is None:
             conversation_config = ConversationAdapters.get_default_conversation_config()
 

src/khoj/database/admin.py

@@ -6,6 +6,7 @@
 from django.http import HttpResponse
 
 from khoj.database.models import (
+    Agent,
     ChatModelOptions,
     ClientApplication,
     Conversation,
@@ -50,6 +51,7 @@ class KhojUserAdmin(UserAdmin):
 admin.site.register(UserSearchModelConfig)
 admin.site.register(TextToImageModelConfig)
 admin.site.register(ClientApplication)
+admin.site.register(Agent)
 
 
 @admin.register(Entry)

src/khoj/database/migrations/0031_agent_conversation_agent.py

@@ -0,0 +1,53 @@
+# Generated by Django 4.2.10 on 2024-03-13 07:38
+
+import django.db.models.deletion
+from django.conf import settings
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0030_conversation_slug_and_title"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="Agent",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("updated_at", models.DateTimeField(auto_now=True)),
+                ("name", models.CharField(max_length=200)),
+                ("tuning", models.TextField()),
+                ("avatar", models.URLField(blank=True, default=None, max_length=400, null=True)),
+                ("tools", models.JSONField(default=list)),
+                ("public", models.BooleanField(default=False)),
+                ("managed_by_admin", models.BooleanField(default=False)),
+                ("slug", models.CharField(blank=True, default=None, max_length=200, null=True)),
+                (
+                    "chat_model",
+                    models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="database.chatmodeloptions"),
+                ),
+                (
+                    "creator",
+                    models.ForeignKey(
+                        blank=True,
+                        default=None,
+                        null=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to=settings.AUTH_USER_MODEL,
+                    ),
+                ),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+        migrations.AddField(
+            model_name="conversation",
+            name="agent",
+            field=models.ForeignKey(
+                blank=True, default=None, null=True, on_delete=django.db.models.deletion.SET_NULL, to="database.agent"
+            ),
+        ),
+    ]

src/khoj/database/models/__init__.py

@@ -1,7 +1,11 @@
 import uuid
+from random import choice
 
 from django.contrib.auth.models import AbstractUser
+from django.core.exceptions import ValidationError
 from django.db import models
+from django.db.models.signals import pre_save
+from django.dispatch import receiver
 from pgvector.django import VectorField
 from phonenumber_field.modelfields import PhoneNumberField
 
@@ -69,6 +73,47 @@ class Type(models.TextChoices):
     renewal_date = models.DateTimeField(null=True, default=None, blank=True)
 
 
+class ChatModelOptions(BaseModel):
+    class ModelType(models.TextChoices):
+        OPENAI = "openai"
+        OFFLINE = "offline"
+
+    max_prompt_size = models.IntegerField(default=None, null=True, blank=True)
+    tokenizer = models.CharField(max_length=200, default=None, null=True, blank=True)
+    chat_model = models.CharField(max_length=200, default="mistral-7b-instruct-v0.1.Q4_0.gguf")
+    model_type = models.CharField(max_length=200, choices=ModelType.choices, default=ModelType.OFFLINE)
+
+
+class Agent(BaseModel):
+    creator = models.ForeignKey(KhojUser, on_delete=models.CASCADE, default=None, null=True, blank=True)
+    name = models.CharField(max_length=200)
+    tuning = models.TextField()
+    avatar = models.URLField(max_length=400, default=None, null=True, blank=True)
+    tools = models.JSONField(default=list)  # List of tools the agent has access to, like online search or notes search
+    public = models.BooleanField(default=False)
+    managed_by_admin = models.BooleanField(default=False)
+    chat_model = models.ForeignKey(ChatModelOptions, on_delete=models.CASCADE)
+    slug = models.CharField(max_length=200, default=None, null=True, blank=True)
+
+
+@receiver(pre_save, sender=Agent)
+def verify_agent(sender, instance, **kwargs):
+    # check if this is a new instance
+    if instance._state.adding:
+        if Agent.objects.filter(name=instance.name, public=True).exists():
+            raise ValidationError(f"A public Agent with the name {instance.name} already exists.")
+        if Agent.objects.filter(name=instance.name, creator=instance.creator).exists():
+            raise ValidationError(f"A private Agent with the name {instance.name} already exists.")
+
+        slug = instance.name.lower().replace(" ", "-")
+        observed_random_numbers = set()
+        while Agent.objects.filter(slug=slug).exists():
+            random_number = choice([i for i in range(0, 10000) if i not in observed_random_numbers])
+            observed_random_numbers.add(random_number)
+            slug = f"{slug}-{random_number}"
+        instance.slug = slug
+
+
 class NotionConfig(BaseModel):
     token = models.CharField(max_length=200)
     user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
@@ -153,17 +198,6 @@ class ModelType(models.TextChoices):
     model_type = models.CharField(max_length=200, choices=ModelType.choices, default=ModelType.OFFLINE)
 
 
-class ChatModelOptions(BaseModel):
-    class ModelType(models.TextChoices):
-        OPENAI = "openai"
-        OFFLINE = "offline"
-
-    max_prompt_size = models.IntegerField(default=None, null=True, blank=True)
-    tokenizer = models.CharField(max_length=200, default=None, null=True, blank=True)
-    chat_model = models.CharField(max_length=200, default="mistral-7b-instruct-v0.1.Q4_0.gguf")
-    model_type = models.CharField(max_length=200, choices=ModelType.choices, default=ModelType.OFFLINE)
-
-
 class UserConversationConfig(BaseModel):
     user = models.OneToOneField(KhojUser, on_delete=models.CASCADE)
     setting = models.ForeignKey(ChatModelOptions, on_delete=models.CASCADE, default=None, null=True, blank=True)
@@ -180,6 +214,7 @@ class Conversation(BaseModel):
     client = models.ForeignKey(ClientApplication, on_delete=models.CASCADE, default=None, null=True, blank=True)
     slug = models.CharField(max_length=200, default=None, null=True, blank=True)
     title = models.CharField(max_length=200, default=None, null=True, blank=True)
+    agent = models.ForeignKey(Agent, on_delete=models.SET_NULL, default=None, null=True, blank=True)
 
 
 class ReflectiveQuestion(BaseModel):