Add support for system settings + switch to chat engine abstraction

docs/overview/design-principles.md

@@ -9,7 +9,7 @@ This is the type of data Docq is designed to work with.
 
 ### Secure Access
 
-With private, organisation data, access it securely becomes paramount.
+With private, organisation data, accessing it securely becomes paramount.
 
 In the age of AI/LLMs, there are many emerging 3rd-party vendors to choose from; However, it is also very easy to unknowingly leak data to these vendors without careful planning on how to adopt these technologies.
 
@@ -21,7 +21,7 @@ It is important for any business to justify costs and benefits for the use of an
 
 - How a business could get started with Docq within minutes;
 - How a business could continue operating Docq with minimal maintenance;
-- And how any Docq user at a business could benefit from it instantly without length training.
+- And how any Docq user at a business could benefit from it instantly without lengthy training.
 
 All of these are reflected in Docq's user interface and user experience design.
 

mkdocs.yaml

@@ -44,6 +44,7 @@ theme:
 
 
 nav:
+  - 'Home': index.md
   - 'Overview':
     - 'Introduction': overview/introduction.md
     - 'Design Principles': overview/design-principles.md

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "docq"
-version = "0.0.1"
+version = "0.0.2"
 description = "Docq"
 authors = ["Docq.AI Team <[email protected]>"]
 maintainers = ["Docq.AI Team <[email protected]>"]

source/docq/manage_settings.py

@@ -5,6 +5,7 @@
 import sqlite3
 from contextlib import closing
 from enum import Enum
+from typing import Any
 
 from .support.store import get_sqlite_system_file, get_sqlite_usage_file
 
@@ -28,6 +29,10 @@ class SystemSettingsKey(Enum):
     MODEL_VENDOR = "model_vendor"
 
 
+class UserSettingsKey(Enum):
+    """User settings keys."""
+
+
 def _get_sqlite_file(user_id: int = None) -> str:
     """Get the sqlite file for the given user."""
     return get_sqlite_usage_file(user_id) if user_id else get_sqlite_system_file()
@@ -63,14 +68,14 @@ def _update_settings(settings: dict, user_id: int = None) -> bool:
         return True
 
 
-def get_system_settings() -> dict:
+def get_system_settings(key: SystemSettingsKey = None) -> Any | None:
     """Get the system settings."""
-    return _get_settings()
+    return _get_settings() if key is None else _get_settings().get(key.value)
 
 
-def get_user_settings(user_id: int) -> dict:
+def get_user_settings(user_id: int, key: UserSettingsKey = None) -> Any | None:
     """Get the user settings."""
-    return _get_settings(user_id)
+    return _get_settings(user_id) if key is None else _get_settings(user_id).get(key.value)
 
 
 def update_system_settings(settings: dict) -> bool:

source/docq/run_queries.py

@@ -90,11 +90,11 @@ def query(input_: str, feature: FeatureKey, space: SpaceKey, spaces: list[SpaceK
 
     data.append(
         (
-            # MESSAGE_TEMPLATE.format(message=response.content)
-            # if is_chat
-            # else
-            # MESSAGE_WITH_SOURCES_TEMPLATE.format(message=response.response, source=response.get_formatted_sources()),
-            response.content if is_chat else response.response,
+            MESSAGE_TEMPLATE.format(message=response.response)
+            if is_chat
+            else MESSAGE_WITH_SOURCES_TEMPLATE.format(
+                message=response.response, source=response.get_formatted_sources()
+            ),
             False,
             datetime.now(),
         )

source/docq/support/llm.py

@@ -3,12 +3,13 @@
 import logging as log
 
 from langchain.chat_models import ChatOpenAI
-from langchain.llms import OpenAI
-from langchain.prompts.chat import (
-    ChatPromptTemplate,
-    HumanMessagePromptTemplate,
-    SystemMessagePromptTemplate,
-)
+
+# from langchain.llms import OpenAI
+# from langchain.prompts.chat import (
+#     ChatPromptTemplate,
+#     HumanMessagePromptTemplate,
+#     SystemMessagePromptTemplate,
+# )
 from langchain.schema import BaseMessage
 from llama_index import (
     GPTListIndex,
@@ -20,20 +21,21 @@
     StorageContext,
     load_index_from_storage,
 )
+from llama_index.chat_engine import CondenseQuestionChatEngine, SimpleChatEngine
 from llama_index.indices.composability import ComposableGraph
 
 from ..domain import SpaceKey
 from .store import get_index_dir, get_upload_dir
 
-PROMPT_CHAT_SYSTEM = """
-You are an AI assistant helping a human to find information.
-Your conversation with the human is recorded in the chat history below.
+# PROMPT_CHAT_SYSTEM = """
+# You are an AI assistant helping a human to find information.
+# Your conversation with the human is recorded in the chat history below.
 
-History:
-"{history}"
-"""
+# History:
+# "{history}"
+# """
 
-PROMPT_CHAT_HUMAN = "{input}"
+# PROMPT_CHAT_HUMAN = "{input}"
 
 PROMPT_QUESTION = """
 You are an AI assistant helping a human to find information in a collection of documents.
@@ -49,16 +51,16 @@
 Assistant:"""
 
 
-def _get_model() -> OpenAI:
-    return OpenAI(temperature=0, model_name="text-davinci-003")
+# def _get_model() -> OpenAI:
+#     return OpenAI(temperature=0, model_name="text-davinci-003")
 
 
 def _get_chat_model() -> ChatOpenAI:
     return ChatOpenAI(temperature=0, model="gpt-3.5-turbo")
 
 
 def _get_llm_predictor() -> LLMPredictor:
-    return LLMPredictor(llm=_get_model())
+    return LLMPredictor(llm=_get_chat_model())
 
 
 def _get_default_storage_context() -> StorageContext:
@@ -97,13 +99,15 @@ def reindex(space: SpaceKey) -> None:
 
 def run_chat(input_: str, history: str) -> BaseMessage:
     """Chat directly with a LLM with history."""
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            SystemMessagePromptTemplate.from_template(PROMPT_CHAT_SYSTEM),
-            HumanMessagePromptTemplate.from_template(PROMPT_CHAT_HUMAN),
-        ]
-    )
-    output = _get_chat_model()(prompt.format_prompt(history=history, input=input_).to_messages())
+    # prompt = ChatPromptTemplate.from_messages(
+    #     [
+    #         SystemMessagePromptTemplate.from_template(PROMPT_CHAT_SYSTEM),
+    #         HumanMessagePromptTemplate.from_template(PROMPT_CHAT_HUMAN),
+    #     ]
+    # )
+    # output = _get_chat_model()(prompt.format_prompt(history=history, input=input_).to_messages())
+    engine = SimpleChatEngine.from_defaults(service_context=_get_service_context())
+    output = engine.chat(input_)
 
     log.debug("(Chat) Q: %s, A: %s", input_, output)
     return output
@@ -125,12 +129,10 @@ def run_ask(input_: str, history: str, space: SpaceKey, spaces: list[SpaceKey] =
         )
         output = graph.as_query_engine().query(PROMPT_QUESTION.format(history=history, input=input_))
     else:
-        # No additional spaces likely to be against a user's documents in their personal space.
+        # No additional spaces i.e. likely to be against a user's documents in their personal space.
         index = _load_index_from_storage(space)
-        output = index.as_query_engine(
-            similarity_top_k=3,
-            vector_store_query_mode="default",
-        ).query(PROMPT_QUESTION.format(history=history, input=input_))
+        engine = index.as_chat_engine(verbose=True, similarity_top_k=3, vector_store_query_mode="default")
+        output = engine.chat(input_)
 
     log.debug("(Ask w/ spaces ) Q: %s, A: %s", spaces, input_, output)
     return output

tests/docq/support/llm_test.py

@@ -11,6 +11,7 @@
 from langchain.schema import BaseMessage
 from llama_index.query_engine.graph_query_engine import ComposableGraphQueryEngine
 from llama_index.indices.composability import ComposableGraph
+from llama_index.chat_engine import SimpleChatEngine
 
 
 # def test_run_ask_with_personal_space_only():
@@ -60,12 +61,18 @@
 
 
 def test_run_chat():
-    with patch("docq.support.llm._get_chat_model") as mock_get_chat_model:
-        mock_chat_openai = Mock(ChatOpenAI)
-        mock_get_chat_model.return_value = mock_chat_openai
-        mock_chat_openai.return_value = "LLM response"
+    with patch.object(SimpleChatEngine, "from_defaults") as mock_simple_chat_engine, patch(
+        "docq.support.llm._get_service_context"
+    ) as mock_get_service_context:
+        mock_get_service_context.return_value = Mock(ServiceContext)
+        mocked_engine = Mock(SimpleChatEngine)
+        mock_simple_chat_engine.return_value = mocked_engine
+        mocked_chat = Mock()
+        mocked_engine.chat = mocked_chat
+        mocked_chat.return_value = "LLM response"
 
         response = run_chat("My ask", "My chat history")
+        mocked_chat.assert_called_once_with("My ask")
         assert response == "LLM response"
 
 

web/personal_ask.py

@@ -3,11 +3,13 @@
 from docq.config import FeatureType
 from docq.domain import FeatureKey
 from st_pages import add_page_title
-from utils.handlers import get_authenticated_user_id
-from utils.layout import auth_required, chat_ui
+from utils.layout import auth_required, chat_ui, feature_enabled
+from utils.sessions import get_authenticated_user_id
 
 auth_required()
 
+feature_enabled(FeatureType.ASK_PERSONAL)
+
 add_page_title()
 
 feature = FeatureKey(FeatureType.ASK_PERSONAL, get_authenticated_user_id())

web/personal_chat.py

@@ -3,11 +3,13 @@
 from docq.config import FeatureType
 from docq.domain import FeatureKey
 from st_pages import add_page_title
-from utils.handlers import get_authenticated_user_id
-from utils.layout import auth_required, chat_ui
+from utils.layout import auth_required, chat_ui, feature_enabled
+from utils.sessions import get_authenticated_user_id
 
 auth_required()
 
+feature_enabled(FeatureType.CHAT_PRIVATE)
+
 add_page_title()
 
 feature = FeatureKey(FeatureType.CHAT_PRIVATE, get_authenticated_user_id())

web/personal_docs.py

@@ -1,13 +1,15 @@
 """Page: Personal / Manage Your Documents."""
 
-from docq.config import SpaceType
+from docq.config import FeatureType, SpaceType
 from docq.domain import SpaceKey
 from st_pages import add_page_title
-from utils.handlers import get_authenticated_user_id
-from utils.layout import auth_required, documents_ui
+from utils.layout import auth_required, documents_ui, feature_enabled
+from utils.sessions import get_authenticated_user_id
 
 auth_required()
 
+feature_enabled(FeatureType.ASK_PERSONAL)
+
 add_page_title()
 
 space = SpaceKey(SpaceType.PERSONAL, get_authenticated_user_id())

web/shared_ask.py

@@ -4,11 +4,13 @@
 from docq.config import FeatureType
 from docq.domain import FeatureKey
 from st_pages import add_page_title
-from utils.handlers import get_authenticated_user_id
-from utils.layout import auth_required, chat_ui, list_spaces_ui
+from utils.layout import auth_required, chat_ui, feature_enabled, list_spaces_ui
+from utils.sessions import get_authenticated_user_id
 
 auth_required()
 
+feature_enabled(FeatureType.ASK_SHARED)
+
 add_page_title()
 
 feature = FeatureKey(FeatureType.ASK_SHARED, get_authenticated_user_id())

web/utils/constants.py

@@ -8,7 +8,9 @@
 class SessionKeySubName(Enum):
     """Second-level names for session keys."""
 
+    CHAT = "chat"
     AUTH = "auth"
+    SETTINGS = "settings"
 
 
 class SessionKeyNameForAuth(Enum):
@@ -19,6 +21,13 @@ class SessionKeyNameForAuth(Enum):
     ADMIN = "admin"
 
 
+class SessionKeyNameForSettings(Enum):
+    """Third-level names for session keys in settings."""
+
+    SYSTEM = "system"
+    USER = "user"
+
+
 class SessionKeyNameForChat(Enum):
     """Third-level names for session keys in chat."""