feat: collection hierarchy

examples/recruiting/candidates_freeform.py

@@ -0,0 +1,42 @@
+# pylint: disable=missing-return-doc, missing-param-doc, missing-function-docstring
+from typing import List
+
+from sqlalchemy import create_engine
+from sqlalchemy.ext.automap import automap_base
+
+from dbally.views.freeform.text2sql import BaseText2SQLView, ColumnConfig, TableConfig
+
+engine = create_engine("sqlite:///examples/recruiting/data/candidates.db")
+
+_Base = automap_base()
+_Base.prepare(autoload_with=engine)
+_Candidate = _Base.classes.candidates
+
+
+class CandidateFreeformView(BaseText2SQLView):
+    """
+    A view for retrieving candidates from the database.
+    """
+
+    def get_tables(self) -> List[TableConfig]:
+        """
+        Get the tables used by the view.
+
+        Returns:
+            A list of tables.
+        """
+        return [
+            TableConfig(
+                name="candidates",
+                columns=[
+                    ColumnConfig("name", "TEXT"),
+                    ColumnConfig("country", "TEXT"),
+                    ColumnConfig("years_of_experience", "INTEGER"),
+                    ColumnConfig("position", "TEXT"),
+                    ColumnConfig("university", "TEXT"),
+                    ColumnConfig("skills", "TEXT"),
+                    ColumnConfig("tags", "TEXT"),
+                    ColumnConfig("id", "INTEGER PRIMARY KEY"),
+                ],
+            ),
+        ]

examples/visualize_fallback_code.py

@@ -0,0 +1,32 @@
+# pylint: disable=missing-function-docstring
+import asyncio
+
+from recruiting import candidate_view_with_similarity_store, candidates_freeform
+from recruiting.candidate_view_with_similarity_store import CandidateView, country_similarity
+from recruiting.candidates_freeform import CandidateFreeformView
+from recruiting.cypher_text2sql_view import SampleText2SQLViewCyphers, create_freeform_memory_engine
+from sqlalchemy import create_engine
+
+import dbally
+from dbally.audit import CLIEventHandler
+from dbally.gradio import create_gradio_interface
+from dbally.llms.litellm import LiteLLM
+
+cm_engine = create_engine("postgresql+pg8000://postgres:ikar89pl@localhost:5432/codebase_community")
+
+
+async def main():
+    await country_similarity.update()
+    llm = LiteLLM(model_name="gpt-3.5-turbo")
+    collection1 = dbally.create_collection("candidates", llm, event_handlers=[CLIEventHandler()])
+    collection2 = dbally.create_collection("freeform candidates", llm, event_handlers=[CLIEventHandler()])
+    collection1.add(CandidateView, lambda: CandidateView(candidate_view_with_similarity_store.engine))
+    collection1.add(SampleText2SQLViewCyphers, lambda: SampleText2SQLViewCyphers(create_freeform_memory_engine()))
+    collection2.add(CandidateFreeformView, lambda: CandidateFreeformView(candidates_freeform.engine))
+    collection1.add_fallback(collection2)
+    gradio_interface = await create_gradio_interface(user_collection=collection1)
+    gradio_interface.launch()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

setup.cfg

@@ -39,6 +39,7 @@ install_requires =
     tabulate>=0.9.0
     click~=8.1.7
     numpy>=1.24.0
+    StrEnum>=0.4.15
 
 [options.extras_require]
 litellm =

src/dbally/_main.py

@@ -14,6 +14,7 @@ def create_collection(
     event_handlers: Optional[List[EventHandler]] = None,
     view_selector: Optional[ViewSelector] = None,
     nl_responder: Optional[NLResponder] = None,
+    fallback_collection: Optional[Collection] = None,
 ) -> Collection:
     """
     Create a new [Collection](collection.md) that is a container for registering views and the\
@@ -44,6 +45,7 @@ def create_collection(
         will be used.
         nl_responder: NL responder used by the collection to respond to natural language queries. If None,\
         a new instance of [NLResponder][dbally.nl_responder.nl_responder.NLResponder] will be used.
+        fallback_collection: Collection to be asked in case of base collection failure.
 
     Returns:
         a new instance of db-ally Collection
@@ -61,4 +63,5 @@ def create_collection(
         view_selector=view_selector,
         llm=llm,
         event_handlers=event_handlers,
+        fallback_collection=fallback_collection,
     )

src/dbally/audit/event_handlers/base.py

@@ -2,12 +2,35 @@
 from abc import ABC
 from typing import Generic, Optional, TypeVar
 
+from strenum import StrEnum
+
 from dbally.audit.events import Event, RequestEnd, RequestStart
 
 RequestCtx = TypeVar("RequestCtx")
 EventCtx = TypeVar("EventCtx")
 
 
+class LogLevel(StrEnum):
+    """
+    An enumeration representing different logging levels.
+
+    This enumeration inherits from `StrEnum`, making each log level a string value.
+    The log levels indicate the severity or type of events that occur within an application,
+    and they are commonly used to filter and categorize log messages.
+
+    Attributes:
+        INFO (str): Represents informational messages that highlight the progress of the application.
+        WARNING (str): Represents potentially harmful situations that require attention.
+        ERROR (str): Represents error events that might still allow the application to continue running.
+        DEBUG (str): Represents detailed debugging messages useful during development and troubleshooting.
+    """
+
+    INFO = "Info"
+    WARNING = "Warning"
+    ERROR = "Error"
+    DEBUG = "Debug"
+
+
 class EventHandler(Generic[RequestCtx, EventCtx], ABC):
     """
     A base class that every custom handler should inherit from
@@ -59,3 +82,13 @@ async def request_end(self, output: RequestEnd, request_context: RequestCtx) ->
             output: The output of the request.
             request_context: Optional context passed from request_start method
         """
+
+    @abc.abstractmethod
+    async def log_message(self, message: str, log_level: LogLevel = LogLevel.INFO) -> None:
+        """
+        Displays the response from the LLM.
+
+        Args:
+            message: db-ally event to be logged with all the details.
+            log_level: log level/importance
+        """

src/dbally/audit/event_handlers/cli_event_handler.py

@@ -3,6 +3,9 @@
 from sys import stdout
 from typing import Optional
 
+from dbally.audit.event_handlers.base import EventHandler, LogLevel
+from dbally.audit.events import Event, LLMEvent, RequestEnd, RequestStart, SimilarityEvent
+
 try:
     from rich import print as pprint
     from rich.console import Console
@@ -14,8 +17,6 @@
     RICH_OUTPUT = False
     pprint = print  # type: ignore
 
-from dbally.audit.event_handlers.base import EventHandler
-from dbally.audit.events import Event, LLMEvent, RequestEnd, RequestStart, SimilarityEvent
 
 _RICH_FORMATING_KEYWORD_SET = {"green", "orange", "grey", "bold", "cyan"}
 _RICH_FORMATING_PATTERN = rf"\[.*({'|'.join(_RICH_FORMATING_KEYWORD_SET)}).*\]"
@@ -98,6 +99,30 @@ async def event_start(self, event: Event, request_context: None) -> None:
                 f"[cyan bold]FETCHER: [grey53]{event.fetcher}\n"
             )
 
+    async def log_message(self, message: str, log_level: LogLevel = LogLevel.INFO) -> None:
+        """
+        Displays message logged by user
+
+        Args:
+            message: Message to be sent
+            log_level: Message log level.
+        """
+        colour = None
+        if log_level == LogLevel.INFO:
+            colour = "white"
+        elif log_level == LogLevel.WARNING:
+            colour = "orange"
+        elif log_level == LogLevel.ERROR:
+            colour = "red"
+        elif log_level == LogLevel.DEBUG:
+            colour = "blue"
+
+        self._print_syntax("[grey53]\n=======================================")
+        self._print_syntax("[grey53]=======================================")
+        self._print_syntax(f"[{colour}]{log_level}: {message}")
+        self._print_syntax("[grey53]=======================================")
+        self._print_syntax("[grey53]=======================================\n")
+
     async def event_end(self, event: Optional[Event], request_context: None, event_context: None) -> None:
         """
         Displays the response from the LLM.
@@ -124,8 +149,11 @@ async def request_end(self, output: RequestEnd, request_context: Optional[dict]
             output: The output of the request.
             request_context: Optional context passed from request_start method
         """
-        self._print_syntax("[green bold]REQUEST OUTPUT:")
-        self._print_syntax(f"Number of rows: {len(output.result.results)}")
+        if output.result:
+            self._print_syntax("[green bold]REQUEST OUTPUT:")
+            self._print_syntax(f"Number of rows: {len(output.result.results)}")
 
-        if "sql" in output.result.context:
-            self._print_syntax(f"{output.result.context['sql']}", "psql")
+            if "sql" in output.result.context:
+                self._print_syntax(f"{output.result.context['sql']}", "psql")
+        else:
+            self._print_syntax("[red bold]No results found")

src/dbally/audit/event_tracker.py

@@ -1,7 +1,7 @@
 from contextlib import asynccontextmanager
 from typing import AsyncIterator, Dict, List, Optional
 
-from dbally.audit.event_handlers.base import EventHandler
+from dbally.audit.event_handlers.base import EventHandler, LogLevel
 from dbally.audit.events import Event, RequestEnd, RequestStart
 from dbally.audit.spans import EventSpan
 
@@ -92,3 +92,15 @@ async def track_event(self, event: Event) -> AsyncIterator[EventSpan]:
             await handler.event_end(
                 span.data, event_context=contexts[handler], request_context=self._request_contexts[handler]
             )
+
+    async def log_message(self, message: str, log_level: LogLevel = LogLevel.INFO) -> None:
+        """
+        Send message to the handler
+
+        Args:
+            message: Message to be sent to
+            log_level: Log level.
+        """
+
+        for handler in self._handlers:
+            await handler.log_message(message, log_level)