fix(assistant): include current time in the taxonomy agent and in the schema generator

ee/hogai/eval/tests/test_eval_funnel_generator.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+from datetime import datetime
 from typing import cast
 
 import pytest
@@ -44,3 +45,17 @@ def test_node_replaces_equals_with_contains(call_node):
     assert "icontains" in actual_output
     assert "John" not in actual_output
     assert "john" in actual_output
+
+
+def test_current_date(call_node):
+    query = "what is the conversion rate from a page view to a next page view in this January?"
+    plan = """Sequence:
+    1. $pageview
+    2. $pageview
+    """
+    date_range = call_node(query, plan).dateRange
+    assert date_range is not None
+    year = str(datetime.now().year)
+    assert (date_range.date_from and year in date_range.date_from) or (
+        date_range.date_to and year in date_range.date_to
+    )

ee/hogai/eval/tests/test_eval_retention_generator.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+from datetime import datetime
 from typing import cast
 
 import pytest
@@ -74,3 +75,19 @@ def test_basic_retention_structure(call_node):
     assert actual_output.retentionFilter.returningEntity == RetentionEntity(
         id="file_uploaded", type="events", name="file_uploaded", order=0
     )
+
+
+def test_current_date(call_node):
+    query = "Show retention for users who signed up in this January?"
+    plan = """Target Event:
+    - signed_up
+
+    Returning Event:
+    - file_uploaded
+    """
+    date_range = call_node(query, plan).dateRange
+    assert date_range is not None
+    year = str(datetime.now().year)
+    assert (date_range.date_from and year in date_range.date_from) or (
+        date_range.date_to and year in date_range.date_to
+    )

ee/hogai/eval/tests/test_eval_trends_generator.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+from datetime import datetime
 from typing import cast
 
 import pytest
@@ -63,3 +64,17 @@ def test_node_leans_towards_line_graph(call_node):
     assert actual_output.series[1].kind == "EventsNode"
     assert actual_output.series[1].event == "downloaded_file"
     assert actual_output.series[1].math == "median_count_per_actor"
+
+
+def test_current_date(call_node):
+    query = "How often do users view the website in this January?"
+    plan = """Events:
+    - $pageview
+        - math operation: total count
+    """
+    date_range = call_node(query, plan).dateRange
+    assert date_range is not None
+    year = str(datetime.now().year)
+    assert (date_range.date_from and year in date_range.date_from) or (
+        date_range.date_to and year in date_range.date_to
+    )

ee/hogai/funnels/prompts.py

@@ -1,6 +1,7 @@
 REACT_SYSTEM_PROMPT = """
 <agent_info>
 You are an expert product analyst agent specializing in data visualization and funnel analysis. Your primary task is to understand a user's data taxonomy and create a plan for building a visualization that answers the user's question. This plan should focus on funnel insights, including a sequence of events, property filters, and values of property filters.
+Current time is {{project_datetime}} in the project's timezone, {{project_timezone}}.
 
 {{core_memory_instructions}}
 </agent_info>
@@ -61,6 +62,7 @@
 
 FUNNEL_SYSTEM_PROMPT = """
 Act as an expert product manager. Your task is to generate a JSON schema of funnel insights. You will be given a generation plan describing a series sequence, filters, exclusion steps, and breakdown. Use the plan and following instructions to create a correct query answering the user's question.
+Current time is {{project_datetime}} in the project's timezone, {{project_timezone}}.
 
 Below is the additional context.
 

ee/hogai/retention/prompts.py

@@ -1,6 +1,7 @@
 REACT_SYSTEM_PROMPT = """
 <agent_info>
 You are an expert product analyst agent specializing in data visualization and retention analysis. Your primary task is to understand a user's data taxonomy and create a plan for building a visualization that answers the user's question. This plan should focus on retention insights, including the target event, returning event, property filters, and values of property filters.
+Current time is {{project_datetime}} in the project's timezone, {{project_timezone}}.
 
 {{core_memory_instructions}}
 </agent_info>
@@ -45,6 +46,7 @@
 
 RETENTION_SYSTEM_PROMPT = """
 Act as an expert product manager. Your task is to generate a JSON schema of retention insights. You will be given a generation plan describing an target event, returning event, target/returning parameters, and filters. Use the plan and following instructions to create a correct query answering the user's question.
+Current time is {{project_datetime}} in the project's timezone, {{project_timezone}}.
 
 Below is the additional context.
 

ee/hogai/schema_generator/nodes.py

@@ -86,7 +86,13 @@ def _run_with_prompt(
         chain = generation_prompt | merger | self._model | parser
 
         try:
-            message: SchemaGeneratorOutput[Q] = chain.invoke({}, config)
+            message: SchemaGeneratorOutput[Q] = chain.invoke(
+                {
+                    "project_datetime": self.project_now,
+                    "project_timezone": self.project_timezone,
+                },
+                config,
+            )
         except PydanticOutputParserException as e:
             # Generation step is expensive. After a second unsuccessful attempt, it's better to send a failure message.
             if len(intermediate_steps) >= 2:

ee/hogai/summarizer/nodes.py

@@ -1,11 +1,9 @@
-import datetime
 import json
 from time import sleep
 from uuid import uuid4
 
 from django.conf import settings
 from django.core.serializers.json import DjangoJSONEncoder
-from django.utils import timezone
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import RunnableConfig
 from langchain_openai import ChatOpenAI
@@ -99,9 +97,6 @@ def run(self, state: AssistantState, config: RunnableConfig) -> PartialAssistant
 
         chain = summarization_prompt | self._model
 
-        utc_now = timezone.now().astimezone(datetime.UTC)
-        project_now = utc_now.astimezone(self._team.timezone_info)
-
         try:
             results = self._compress_results(viz_message, results_response["results"])
             example_prompt = self._get_example_prompt(viz_message)
@@ -118,9 +113,9 @@ def run(self, state: AssistantState, config: RunnableConfig) -> PartialAssistant
                 "query_kind": viz_message.answer.kind,
                 "core_memory": self.core_memory_text,
                 "results": results,
-                "utc_datetime_display": utc_now.strftime("%Y-%m-%d %H:%M:%S"),
-                "project_datetime_display": project_now.strftime("%Y-%m-%d %H:%M:%S"),
-                "project_timezone": self._team.timezone_info.tzname(utc_now),
+                "utc_datetime_display": self.utc_now,
+                "project_datetime_display": self.project_now,
+                "project_timezone": self.project_timezone,
                 "example": example_prompt,
             },
             config,
@@ -151,7 +146,7 @@ def _compress_results(self, viz_message: VisualizationMessage, results: list[dic
             return compress_and_format_trends_results(results)
         elif isinstance(viz_message.answer, AssistantFunnelsQuery):
             query_date_range = QueryDateRange(
-                viz_message.answer.dateRange, self._team, viz_message.answer.interval, datetime.datetime.now()
+                viz_message.answer.dateRange, self._team, viz_message.answer.interval, self._utc_now_datetime
             )
             funnel_step_reference = (
                 viz_message.answer.funnelsFilter.funnelStepReference if viz_message.answer.funnelsFilter else None

ee/hogai/taxonomy_agent/nodes.py

@@ -97,6 +97,8 @@ def _run_with_prompt_and_toolkit(
                         "events": self._events_prompt,
                         "agent_scratchpad": self._get_agent_scratchpad(intermediate_steps),
                         "core_memory_instructions": CORE_MEMORY_INSTRUCTIONS,
+                        "project_datetime": self.project_now,
+                        "project_timezone": self.project_timezone,
                     },
                     config,
                 ),

ee/hogai/trends/prompts.py

@@ -1,6 +1,7 @@
 REACT_SYSTEM_PROMPT = """
 <agent_info>
 You are an expert product analyst agent specializing in data visualization and trends analysis. Your primary task is to understand a user's data taxonomy and create a plan for building a visualization that answers the user's question. This plan should focus on trends insights, including a series of events, property filters, and values of property filters.
+Current time is {{project_datetime}} in the project's timezone, {{project_timezone}}.
 
 {{core_memory_instructions}}
 </agent_info>
@@ -105,6 +106,7 @@
 
 TRENDS_SYSTEM_PROMPT = """
 Act as an expert product manager. Your task is to generate a JSON schema of trends insights. You will be given a generation plan describing series, filters, and breakdowns. Use the plan and following instructions to create a correct query answering the user's question.
+Current time is {{project_datetime}} in the project's timezone, {{project_timezone}}.
 
 Below is the additional context.
 

ee/hogai/utils/nodes.py

@@ -1,5 +1,7 @@
+import datetime
 from abc import ABC, abstractmethod
 
+from django.utils import timezone
 from langchain_core.runnables import RunnableConfig
 
 from ee.models.assistant import CoreMemory
@@ -30,3 +32,28 @@ def core_memory_text(self) -> str:
         if not self.core_memory:
             return ""
         return self.core_memory.formatted_text
+
+    @property
+    def _utc_now_datetime(self) -> datetime.datetime:
+        return timezone.now().astimezone(datetime.UTC)
+
+    @property
+    def utc_now(self) -> str:
+        """
+        Returns the current time in UTC.
+        """
+        return self._utc_now_datetime.strftime("%Y-%m-%d %H:%M:%S")
+
+    @property
+    def project_now(self) -> str:
+        """
+        Returns the current time in the project's timezone.
+        """
+        return self._utc_now_datetime.astimezone(self._team.timezone_info).strftime("%Y-%m-%d %H:%M:%S")
+
+    @property
+    def project_timezone(self) -> str | None:
+        """
+        Returns the timezone of the project, e.g. "PST" or "UTC".
+        """
+        return self._team.timezone_info.tzname(self._utc_now_datetime)