ohdearquant · September 6, 2024 22:32 · Sep 6, 2024 · Sep 6, 2024
diff --git a/stock_lookup_assistant.py b/stock_lookup_assistant.py
@@ -20,16 +20,6 @@ class Speaker(str, Enum):
 
 chat_model = iModel(**model_config)
 
-stock_system_prompt = ("""
-       
-    You can only look up stock symbols given to you by the search_for_stock_symbol tool, don't make them up. Trust the output of the search_for_stock_symbol tool even if it doesn't make sense to you.
-    The current user state is:
-    {state}
-    Once you have supplied a stock price, you must call the tool "done" to signal that you are done.
-    If the user asks to do anything other than look up a stock symbol or price, mark the is_done flag"
-    to signal some other agent should help.
-""")
-
 
 class StockLookupForm(Form):
 

diff --git a/stock_lookup_assistant.py b/stock_lookup_assistant.py
@@ -0,0 +1,111 @@
+from enum import Enum
+import pprint
+from typing import ClassVar
+from pydantic import Field
+
+from lionagi import Branch, iModel, pile, Form
+
+class Speaker(str, Enum):
+    STOCK_LOOKUP = "stock_lookup"
+    AUTHENTICATE = "authenticate"
+    ACCOUNT_BALANCE = "account_balance"
+    TRANSFER_MONEY = "transfer_money"
+    CONCIERGE = "concierge"
+    ORCHESTRATOR = "orchestrator"
+
+model_config = {
+    "model": "gpt-4o-mini",
+    "temperature": 0.3,
+}
+
+chat_model = iModel(**model_config)
+
+stock_system_prompt = ("""
+       
+    You can only look up stock symbols given to you by the search_for_stock_symbol tool, don't make them up. Trust the output of the search_for_stock_symbol tool even if it doesn't make sense to you.
+    The current user state is:
+    {state}
+    Once you have supplied a stock price, you must call the tool "done" to signal that you are done.
+    If the user asks to do anything other than look up a stock symbol or price, mark the is_done flag"
+    to signal some other agent should help.
+""")
+
+
+class StockLookupForm(Form):
+
+    system_prompt: ClassVar = "You are a helpful assistant that is looking up stock prices."
+
+    instruction_prompt: ClassVar = "The user may not know the stock symbol of the company they're interested in,"
+    "so you can help them look it up by the name of the company. If the user asks to do anything other than "
+    "look up a stock symbol or price, set stock_price value to be -1"
+
+    context: str | dict | None = Field(
+        None,
+        title="Context",
+        description="The context of the task.",
+    )
+
+    company_name: str | None = Field(
+        None,
+        title="Company Name",
+        description="The name of the company user is interested in.",
+    )
+
+    stock_ticker: str | None = Field(
+        None,
+        title="Stock Ticker",
+        description="The stock ticker of the company user is interested in.",
+    )
+
+    stock_price: float | int | None = Field(
+        None,
+        title="Stock Price",
+        description="The price of the stock.",
+    )
+
+    is_done: bool = False
+
+    assignment: str = "task -> stock_price"
+
+    def __init__(
+        self,
+        instruction: str = None,
+        context: str | dict = None,
+        company_name: str | None = None,
+        stock_ticker: str | None = None, 
+    ):
+        super().__init__()
+        task = f"""
+        You are tasked with looking up stock prices.
+        - additional instruction: {instruction}.
+        - additional context: {context}.
+        """
+
+        self.task = task
+
+        if company_name:
+            self.append_to_input("company_name", company_name)
+
+        if stock_ticker:
+            self.append_to_input("stock_ticker", stock_ticker)
+
+    @staticmethod
+    def lookup_stock_price(stock_symbol: str) -> float:
+        """
+        Useful for looking up a stock price. You can only look up stock 
+        symbols given to you by the search_for_stock_symbol tool,
+        don't make them up.
+        """
+        print(f"Looking up stock price for {stock_symbol}")
+        return 100.0
+
+    @staticmethod
+    def search_for_stock_symbol(str_: str) -> str:
+        """
+        Useful for searching for a stock symbol given a free-form 
+        company name. Trust the output of the search_for_stock_symbol 
+        tool even if it doesn't make sense to you
+        """
+        print("Searching for stock symbol")
+        return str_.upper()
+