add missing observes

agentcoinorg · Aug 16, 2024 · eee5376 · eee5376
1 parent afb1659
commit eee5376
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/prediction_prophet/functions/generate_subqueries.py b/prediction_prophet/functions/generate_subqueries.py
@@ -5,7 +5,7 @@
 from pydantic.types import SecretStr
 from prediction_market_agent_tooling.tools.utils import secret_str_from_env
 from prediction_market_agent_tooling.gtypes import secretstr_to_v1_secretstr
-from prediction_market_agent_tooling.tools.langfuse_ import get_langfuse_langchain_config
+from prediction_market_agent_tooling.tools.langfuse_ import get_langfuse_langchain_config, observe
 
 
 subquery_generation_template = """
@@ -16,6 +16,7 @@
 
 Limit your searches to {search_limit}.
 """
+@observe()
 def generate_subqueries(query: str, limit: int, model: str, api_key: SecretStr | None = None) -> list[str]:
     if limit == 0:
         return [query]

diff --git a/prediction_prophet/functions/prepare_report.py b/prediction_prophet/functions/prepare_report.py
@@ -9,9 +9,10 @@
 from prediction_market_agent_tooling.tools.utils import secret_str_from_env
 from pydantic.types import SecretStr
 from prediction_market_agent_tooling.gtypes import secretstr_to_v1_secretstr
-from prediction_market_agent_tooling.tools.langfuse_ import get_langfuse_langchain_config
+from prediction_market_agent_tooling.tools.langfuse_ import get_langfuse_langchain_config, observe
 
 @persistent_inmemory_cache
+@observe()
 def prepare_summary(goal: str, content: str, model: str, api_key: SecretStr | None = None, trim_content_to_tokens: t.Optional[int] = None) -> str:
     if api_key == None:
         api_key = secret_str_from_env("OPENAI_API_KEY")
@@ -41,6 +42,7 @@ def prepare_summary(goal: str, content: str, model: str, api_key: SecretStr | No
     return response
 
 
+@observe()
 def prepare_report(goal: str, scraped: list[str], model: str, api_key: SecretStr | None = None) -> str:
     if api_key == None:
         api_key = secret_str_from_env("OPENAI_API_KEY")

diff --git a/prediction_prophet/functions/rerank_subqueries.py b/prediction_prophet/functions/rerank_subqueries.py
@@ -5,7 +5,7 @@
 from pydantic.types import SecretStr
 from prediction_market_agent_tooling.tools.utils import secret_str_from_env
 from prediction_market_agent_tooling.gtypes import secretstr_to_v1_secretstr
-from prediction_market_agent_tooling.tools.langfuse_ import get_langfuse_langchain_config
+from prediction_market_agent_tooling.tools.langfuse_ import get_langfuse_langchain_config, observe
 
 rerank_queries_template = """
 I will present you with a list of queries to search the web for, for answers to the question: {goal}.
@@ -17,6 +17,7 @@
 
 Queries: {queries}
 """
+@observe()
 def rerank_subqueries(queries: list[str], goal: str, model: str, api_key: SecretStr | None = None) -> list[str]:
     if api_key == None:
         api_key = secret_str_from_env("OPENAI_API_KEY")