Add additional context (#83)

* Add `Session.add_system_instructions` method

* Rename `Session.add_system_instructions` to `Session.add_additional_context`

* Add basic add_additional_context unit tests

* Allow add_additional_context to be called multiple times

* Address PR comments

* Address PR comments

* Clearer properties

* Address PR comments

Author: mare5x

databao/agents/lighthouse/agent.py

@@ -17,24 +17,29 @@ def __init__(self) -> None:
         """Initialize agent with lazy graph compilation."""
         super().__init__()
         self._cached_graph: ExecuteSubmit | None = None
+        self._prompt_template = read_prompt_template(Path("system_prompt.jinja"))
 
     def render_system_prompt(self, data_connection: Any, session: Session) -> str:
         """Render system prompt with database schema."""
-        prompt_template = read_prompt_template(Path("system_prompt.jinja"))
         db_schema = describe_duckdb_schema(data_connection)
-        db_contexts, df_contexts = session.context
+
         context = ""
-        for db_name, db_context in db_contexts.items():
+        for db_name, db_context in session.db_context.items():
             context += f"## Context for DB {db_name}\n\n{db_context}\n\n"
-        for df_name, df_context in df_contexts.items():
+        for df_name, df_context in session.df_context.items():
             context += f"## Context for DF {df_name} (fully qualified name 'temp.main.{df_name}')\n\n{df_context}\n\n"
+        for idx, additional_ctx in enumerate(session.additional_context, start=1):
+            additional_context = additional_ctx.strip()
+            context += f"## General information {idx}\n\n{additional_context}\n\n"
+        context = context.strip()
 
-        prompt = prompt_template.render(
+        prompt = self._prompt_template.render(
             date=get_today_date_str(),
             db_schema=db_schema,
             context=context,
         )
-        return prompt
+
+        return prompt.strip()
 
     def _create_graph(self, data_connection: Any, llm_config: LLMConfig) -> Any:
         """Create and compile the Lighthouse agent graph."""

databao/core/session.py

@@ -40,8 +40,9 @@ def __init__(
         self.__dbs: dict[str, Any] = {}
         self.__dfs: dict[str, DataFrame] = {}
 
-        self.__db_contexts: dict[str, str] = {}
-        self.__df_contexts: dict[str, str] = {}
+        self.__db_context: dict[str, str] = {}
+        self.__df_context: dict[str, str] = {}
+        self.__additional_context: list[str] = []
 
         # Create a DuckDB connection for the session
         self.__duckdb_connection = duckdb.connect(":memory:")
@@ -93,7 +94,7 @@ def add_db(self, connection: Any, *, name: str | None = None, context: str | Pat
             self.__dbs[conn_name] = connection
 
         if (context_text := self._parse_context_arg(context)) is not None:
-            self.__db_contexts[conn_name] = context_text
+            self.__db_context[conn_name] = context_text
 
     def add_df(self, df: DataFrame, *, name: str | None = None, context: str | Path | None = None) -> None:
         """Register a DataFrame in this session and in the session's DuckDB.
@@ -114,7 +115,21 @@ def add_df(self, df: DataFrame, *, name: str | None = None, context: str | Path
             self.__dbs["duckdb"] = self.__duckdb_connection
 
         if (context_text := self._parse_context_arg(context)) is not None:
-            self.__df_contexts[df_name] = context_text
+            self.__df_context[df_name] = context_text
+
+    def add_context(self, context: str | Path) -> None:
+        """Add additional context to help models understand your data.
+
+        Use this method to add general information that might not be associated with a specific data source.
+        If the information is specific to a data source, use the `context` argument of `add_db` and `add_df`.
+
+        Args:
+            context: The string or the path to a file containing the additional context.
+        """
+        text = self._parse_context_arg(context)
+        if text is None:
+            raise ValueError("Invalid context provided.")
+        self.__additional_context.append(text)
 
     def thread(self, *, stream_ask: bool | None = None, stream_plot: bool | None = None) -> Pipe:
         """Start a new thread in this session."""
@@ -158,6 +173,16 @@ def cache(self) -> "Cache":
         return self.__cache
 
     @property
-    def context(self) -> tuple[dict[str, str], dict[str, str]]:
-        """Per-source natural-language context for DBs and DFs: (db_contexts, df_contexts)."""
-        return self.__db_contexts, self.__df_contexts
+    def db_context(self) -> dict[str, str]:
+        """Per-source natural-language context for DBs."""
+        return self.__db_context
+
+    @property
+    def df_context(self) -> dict[str, str]:
+        """Per-source natural-language context for DFs."""
+        return self.__df_context
+
+    @property
+    def additional_context(self) -> list[str]:
+        """General additional context not specific to any one data source."""
+        return self.__additional_context

tests/test_session.py

@@ -41,9 +41,8 @@ def test_add_db_with_temp_file_context(temp_context_file: Path) -> None:
     session = databao.open_session("ctx_path_session_db_file")
     session.add_db(conn, context=temp_context_file)
 
-    db_contexts, _df_contexts = session.context
-    assert "db1" in db_contexts
-    assert db_contexts["db1"] == temp_context_file.read_text()
+    assert "db1" in session.db_context
+    assert session.db_context["db1"] == temp_context_file.read_text()
 
 
 def test_add_df_with_temp_file_context(temp_context_file: Path) -> None:
@@ -52,9 +51,8 @@ def test_add_df_with_temp_file_context(temp_context_file: Path) -> None:
     session = databao.open_session("ctx_path_session_df_file")
     session.add_df(df, context=temp_context_file)
 
-    _db_contexts, df_contexts = session.context
-    assert "df1" in df_contexts
-    assert df_contexts["df1"] == temp_context_file.read_text()
+    assert "df1" in session.df_context
+    assert session.df_context["df1"] == temp_context_file.read_text()
 
 
 def test_add_db_with_string_context() -> None:
@@ -64,9 +62,8 @@ def test_add_db_with_string_context() -> None:
     context_string = "This is a string context for the database."
     session.add_db(conn, context=context_string)
 
-    db_contexts, _df_contexts = session.context
-    assert "db1" in db_contexts
-    assert db_contexts["db1"] == context_string
+    assert "db1" in session.db_context
+    assert session.db_context["db1"] == context_string
 
 
 def test_add_df_with_string_context() -> None:
@@ -76,6 +73,43 @@ def test_add_df_with_string_context() -> None:
     context_string = "This is a string context for the DataFrame."
     session.add_df(df, context=context_string)
 
-    _db_contexts, df_contexts = session.context
-    assert "df1" in df_contexts
-    assert df_contexts["df1"] == context_string
+    assert "df1" in session.df_context
+    assert session.df_context["df1"] == context_string
+
+
+def test_add_additional_context_with_nonexistent_path_raises() -> None:
+    """add_additional_context should raise if given a non-existent Path."""
+    session = databao.open_session("additional_ctx_missing_path")
+    with pytest.raises(FileNotFoundError):
+        session.add_context(Path("no_such_context_file_123.md"))
+
+
+def test_add_additional_context_with_temp_file(temp_context_file: Path) -> None:
+    """Ensure additional context can be loaded from a temporary file path."""
+    session = databao.open_session("additional_ctx_from_file")
+    session.add_context(temp_context_file)
+    assert session.additional_context == [temp_context_file.read_text()]
+
+
+def test_add_additional_context_with_string() -> None:
+    """Ensure additional context can be provided directly as a string."""
+    session = databao.open_session("additional_ctx_from_string")
+    text = "Global instructions for the session go here."
+    session.add_context(text)
+    assert session.additional_context == [text]
+
+
+def test_add_additional_context_multiple_calls_mixed_sources(temp_context_file: Path) -> None:
+    """Calling add_additional_context multiple times should append in order."""
+    session = databao.open_session("additional_ctx_multiple")
+    first = "First global instruction."
+    second = temp_context_file.read_text()
+    third = "Third bit of context."
+
+    session.add_context(first)
+    session.add_context(temp_context_file)
+    session.add_context(third)
+
+    assert first in session.additional_context
+    assert second in session.additional_context
+    assert third in session.additional_context