Feat/assistant (#98)

* 🚧 Wip(assistant): checkpoint stash before vacay

working on modularizing and utilizing "assistant templates" - for now a rewrite of LLM components
(will merge/replace with llm/ in final)

* ✨ Feat(assistant): assistant api working with rag

* ✨ Feat(assistant): change chat to use assistants

* ♻️  Refactoring(assistant): move and cleanup old chat code

* ✨ Feat(assistant): crud ops + move initial to json files

* 🎨 Style(assistant): minor comment fix

.gitignore

@@ -149,6 +149,6 @@ Pipfile.lock
 pyrightconfig.json
 
 # Project specific ignores (add your own here)
-data/
+src/planning_permission/data/
 output/
 *.xlsx

data/assistant-templates/example-multi.json

@@ -0,0 +1,39 @@
+{
+  "id": "multi_example",
+  "name": "Multi-stream Example",
+  "description": "This is an example assistant that showcases how an assistant can run through multiple different internal prompts before answering the user.",
+  "streams": [
+    {
+      "name": "First Stream",
+      "settings": {
+        "model": "gpt-4o"
+      },
+      "messages": [
+        {
+          "role": "system",
+          "content": "Make this text sound more fancy and verbose."
+        },
+        {
+          "role": "user",
+          "content": "{query}"
+        }
+      ]
+    },
+    {
+      "name": "Second Stream",
+      "settings": {
+        "model": "gpt-4o"
+      },
+      "messages": [
+        {
+          "role": "system",
+          "content": "Repeat back any text verbatim and count the number of words"
+        },
+        {
+          "role": "user",
+          "content": "{last_input}"
+        }
+      ]
+    }
+  ]
+}

data/assistant-templates/example-rag.json

@@ -0,0 +1,24 @@
+{
+  "id": "rag_example",
+  "name": "RAG (document) Example",
+  "files_collection_id": "Upload files in the Assistant Creation UI",
+  "description": "This is an example assistant that showcases how to setup an assistant for answering questions about uploaded documents.  \n\nMake sure to upload documents first.",
+  "streams": [
+    {
+      "name": "RagStream",
+      "settings": {
+        "model": "gpt-4o"
+      },
+      "messages": [
+        {
+          "role": "system",
+          "content": "Explain what this is and repeat back an excerpt of it."
+        },
+        {
+          "role": "user",
+          "content": "{rag_results}"
+        }
+      ]
+    }
+  ]
+}

data/assistant-templates/example.json

@@ -0,0 +1,28 @@
+{
+  "id": "example",
+  "name": "Example Assistant",
+  "description": "This is an *example assistant* for testing purposes.",
+  "sample_questions": [
+    "Hur kan jag kolla upp tidstabellen för bussar?",
+    "Vilka evenemang händer i sommar?",
+    "Var kan jag parkera?"
+  ],
+  "streams": [
+    {
+      "name": "ChatStream",
+      "settings": {
+        "model": "gpt-4o"
+      },
+      "messages": [
+        {
+          "role": "system",
+          "content": "You are a helpful AI assistant that helps people with answering questions related to municipality and Helsingborg City. The questions are going to be asked in Swedish. Your response must always be in Swedish."
+        },
+        {
+          "role": "user",
+          "content": "{query}"
+        }
+      ]
+    }
+  ]
+}

data/assistant-templates/planning-permission.json

@@ -0,0 +1,28 @@
+{
+  "id": "planning",
+  "name": "Planning Assistant",
+  "files_collection_id": "Upload files in the Assistant Creation UI",
+  "description": "Den här assistenten är expert på att svara på frågor angående bygglov i Helsingborg Stad.",
+  "streams": [
+    {
+      "name": "ChatStream",
+      "settings": {
+        "model": "gpt-4o"
+      },
+      "messages": [
+        {
+          "role": "system",
+          "content": "You are a helpful AI assistant that helps people with answering questions about planning permission.<br> If you can't find the answer in the search result below, just say (in Swedish) \"Tyvärr kan jag inte svara på det.\" Don't try to make up an answer.<br> If the question is not related to the context, politely respond that you are tuned to only answer questions that are related to the context.<br> The questions are going to be asked in Swedish. Your response must always be in Swedish."
+        },
+        {
+          "role": "user",
+          "content": "{query}"
+        },
+        {
+          "role": "user",
+          "content": "Here are the results of the search:\n\n {rag_results}"
+        }
+      ]
+    }
+  ]
+}

fai-rag-app/fai-backend/fai_backend/assistant/assistant.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from langstream import Stream
+
+from fai_backend.assistant.models import AssistantTemplate
+from fai_backend.assistant.protocol import ILLMProtocol, IAssistantStreamProtocol
+from fai_backend.llm.service import create_rag_stream
+
+
+class AssistantLLM(ILLMProtocol):
+    def __init__(self, template: AssistantTemplate, base_stream: IAssistantStreamProtocol):
+        self.template = template
+        self.base_stream = base_stream
+        self.vars = {}
+
+    async def create(self) -> Stream[str, Any]:
+        all_stream_producers = [self.base_stream.create_stream(stream_def, lambda: self.vars) for
+                                stream_def in self.template.streams]
+
+        chained_stream: Stream[str, Any] | None = None
+
+        if self.template.files_collection_id:
+            chained_stream = await self._create_rag_stream()
+
+        for stream_producer in all_stream_producers:
+            chained_stream = chained_stream.and_then(await stream_producer) \
+                if chained_stream is not None \
+                else await stream_producer
+
+        def preprocess(initial_query: str):
+            self.vars.update({"query": initial_query})
+            return initial_query
+
+        return (
+            Stream[str, str]('preprocess', preprocess)
+            .and_then(chained_stream)
+        )
+
+    async def _create_rag_stream(self) -> Stream[str, str]:
+        async def run_rag_stream(initial_query: list[str]):
+            stream = await create_rag_stream(initial_query[0], self.template.files_collection_id)
+            async for r in stream(initial_query[0]):
+                yield r
+
+        def rag_postprocess(in_data: Any):
+            results = in_data[0]['results']
+            self.vars.update({'rag_results': results})
+            return self.vars['query']
+
+        return (
+            Stream('RAGStream', run_rag_stream)
+            .and_then(rag_postprocess)
+        )

fai-rag-app/fai-backend/fai_backend/assistant/assistant_openai.py

@@ -0,0 +1,40 @@
+from typing import Callable, Any, Iterable
+
+from langstream import Stream
+from langstream.contrib import OpenAIChatStream, OpenAIChatDelta, OpenAIChatMessage
+
+from fai_backend.assistant.models import LLMStreamDef, LLMStreamMessage
+from fai_backend.assistant.protocol import IAssistantStreamProtocol
+
+
+class OpenAIAssistantStream(IAssistantStreamProtocol):
+    async def create_stream(self, stream_def: LLMStreamDef, get_vars: Callable[[], dict]) -> Stream[str, str]:
+        return OpenAIChatStream[str, OpenAIChatDelta](
+            stream_def.name,
+            lambda in_data: self._to_openai_messages(in_data, stream_def.messages, get_vars),
+            model=stream_def.settings.model,
+            temperature=getattr(stream_def, 'settings.temperature', 0),
+            functions=getattr(stream_def, 'functions', None),
+            function_call=getattr(stream_def, 'function_call', None),
+        ).map(lambda delta: delta.content)
+
+    @staticmethod
+    def _to_openai_messages(
+            in_data: Any,
+            messages: list[LLMStreamMessage],
+            get_vars: Callable[[], dict]
+    ) -> Iterable[OpenAIChatMessage]:
+        def parse_in_data(data: Any):
+            if isinstance(data, list):
+                return "".join([parse_in_data(c) for c in data])
+            return str(data)
+
+        in_data_as_str = parse_in_data(in_data)
+
+        input_vars = get_vars()
+        input_vars['last_input'] = in_data_as_str
+
+        return [OpenAIChatMessage(
+            content=message.content.format(**input_vars),
+            role=message.role
+        ) for message in messages]

fai-rag-app/fai-backend/fai_backend/assistant/models.py

@@ -0,0 +1,37 @@
+from datetime import datetime
+from typing import Optional, Dict, List, Any, Literal, Union
+
+from pydantic import BaseModel
+
+
+class LLMStreamMessage(BaseModel):
+    role: Literal["system", "user", "assistant", "function"]
+    content: str
+
+
+class LLMStreamSettings(BaseModel):
+    model: str
+    temperature: Optional[float] = 0
+    functions: Optional[List[Dict[str, Any]]] = None
+    function_call: Optional[Union[Literal["none", "auto"], Dict[str, Any]]] = None
+
+
+class LLMStreamDef(BaseModel):
+    name: str
+    settings: LLMStreamSettings
+    messages: Optional[List[LLMStreamMessage]] = None
+
+
+class AssistantTemplate(BaseModel):
+    id: str
+    name: str
+    files_collection_id: Optional[str] = None
+    description: Optional[str] = None
+    sample_questions: list[str] = []
+    streams: List[LLMStreamDef]
+
+
+class LLMClientChatMessage(BaseModel):
+    date: datetime
+    source: str | None = None
+    content: str | None = None

fai-rag-app/fai-backend/fai_backend/assistant/protocol.py

@@ -0,0 +1,19 @@
+from typing import Protocol, Callable
+
+from langstream import Stream
+
+from fai_backend.assistant.models import LLMStreamDef
+
+
+class ILLMProtocol(Protocol):
+    async def create(self) -> Stream[str, str]:
+        """
+        Create a Stream that takes a str (generally a question) and returns
+        a stream of tokens (strings) of the response given by the LLM.
+        """
+        ...
+
+
+class IAssistantStreamProtocol(Protocol):
+    async def create_stream(self, stream_def: LLMStreamDef, get_vars: Callable[[], dict]) -> Stream[str, str]:
+        ...

fai-rag-app/fai-backend/fai_backend/assistant/routes.py

@@ -0,0 +1,89 @@
+from fastapi import APIRouter, Depends, Security
+from langstream import join_final_output
+
+from fai_backend.assistant.models import AssistantTemplate
+from fai_backend.assistant.service import AssistantFactory
+from fai_backend.dependencies import get_authenticated_user
+from fai_backend.logger.route_class import APIRouter as LoggingAPIRouter
+from fai_backend.projects.dependencies import list_projects_request, get_project_request, get_project_service, \
+    update_project_request
+from fai_backend.projects.schema import ProjectResponse, ProjectUpdateRequest
+from fai_backend.projects.service import ProjectService
+
+router = APIRouter(
+    prefix='/api',
+    tags=['Assistant'],
+    route_class=LoggingAPIRouter,
+    dependencies=[],
+)
+
+
+@router.get(
+    '/assistant/{project_id}/ask/{assistant_id}',
+    summary="Ask an assistant a question.",
+    dependencies=[Security(get_authenticated_user)]
+)
+async def ask_assistant(
+        project_id: str,
+        assistant_id: str,
+        question: str,
+        projects: list[ProjectResponse] = Depends(list_projects_request),
+):
+    print(f"Assistant: {project_id}/{assistant_id} - {question}")
+    factory = AssistantFactory([a for p in projects for a in p.assistants if p.id == project_id])
+    assistant = factory.create_assistant_stream(assistant_id)
+    stream = await assistant.create()
+    return await join_final_output(stream(question))
+
+
+@router.get(
+    '/assistant/{project_id}/template',
+    summary="Get assistant templates.",
+    response_model=list[AssistantTemplate],
+    dependencies=[Security(get_authenticated_user)]
+)
+async def get_template(
+        project_id: str,
+        projects: list[ProjectResponse] = Depends(list_projects_request)
+):
+    return [a for p in projects for a in p.assistants if p.id == project_id]
+
+
+@router.post(
+    '/assistant/{project_id}/template',
+    summary="Create/update assistant template.",
+    response_model=AssistantTemplate,
+    dependencies=[Security(get_authenticated_user)]
+)
+async def create_template(
+        template: AssistantTemplate,
+        existing_project: ProjectResponse = Depends(get_project_request),
+        project_service: ProjectService = Depends(get_project_service),
+):
+    existing_project.assistants = [a for a in existing_project.assistants if a.id != template.id]
+    existing_project.assistants.append(template)
+    await update_project_request(
+        body=ProjectUpdateRequest(**existing_project.model_dump()),
+        existing_project=existing_project,
+        project_service=project_service)
+    return template
+
+
+@router.delete(
+    '/assistant/{project_id}/template/{assistant_id}',
+    summary="Delete assistant template.",
+    response_model=AssistantTemplate | None,
+    dependencies=[Security(get_authenticated_user)]
+)
+async def delete_template(
+        assistant_id: str,
+        existing_project: ProjectResponse = Depends(get_project_request),
+        project_service: ProjectService = Depends(get_project_service)
+):
+    assistant = next((a for a in existing_project.assistants if a.id == assistant_id), None)
+    existing_project.assistants = [a for a in existing_project.assistants if a.id != assistant_id]
+    await update_project_request(
+        body=ProjectUpdateRequest(**existing_project.model_dump()),
+        existing_project=existing_project,
+        project_service=project_service)
+    return assistant

fai-rag-app/fai-backend/fai_backend/assistant/service.py

@@ -0,0 +1,20 @@
+from fai_backend.assistant.assistant import AssistantLLM
+from fai_backend.assistant.assistant_openai import OpenAIAssistantStream
+from fai_backend.assistant.models import AssistantTemplate
+from fai_backend.config import settings
+from fai_backend.assistant.protocol import IAssistantStreamProtocol
+
+
+class AssistantFactory:
+    def __init__(self, assistant_templates: list[AssistantTemplate]):
+        self.assistant_templates = assistant_templates
+
+    def create_assistant_stream(self, assistant_id: str, backend: str = settings.LLM_BACKEND) -> AssistantLLM:
+        assistant = next(a for a in self.assistant_templates if a.id == assistant_id)
+        return AssistantLLM(assistant, self._get_stream_constructor(backend))
+
+    @staticmethod
+    def _get_stream_constructor(backend: str) -> IAssistantStreamProtocol:
+        return {
+            'openai': lambda: OpenAIAssistantStream(),
+        }[backend]()