Fix google search example

examples/google_search.py

@@ -1,40 +1,18 @@
-import os
-import time
-from concurrent.futures import ThreadPoolExecutor
+import json
 from typing import List, Tuple
 
-import serpapi
-from openai import OpenAI
-
-from memgpt.credentials import MemGPTCredentials
-from memgpt.data_sources.connectors import WebConnector
-from memgpt.utils import printd
+from memgpt import create_client
+from memgpt.agent import Agent
+from memgpt.memory import ChatMemory
 
 """
 This example show how you can add a google search custom function to your MemGPT agent.
 
-First, make sure you run `pip install serpapi`, then setup memgpt:
-
-1. Copy this file into the `~/.memgpt/functions` directory:
-```
-cp examples/google_search.py  ~/.memgpt/functions/google_search.py
-```
-
-2. Create a preset file that include the function `google_search`
-
-3. Add the preset file via the CLI:
-```
-memgpt add preset -f  examples/google_search_preset.yaml --name search_preset
-```
-
-4. Run memgpt with the `google_search_persona` persona:
-```
-memgpt run --preset search_preset --persona google_search_persona
-```
+First, make sure you run `pip install serpapi`, then setup memgpt with `memgpt configure`.
 """
 
 
-def google_search(self, query: str) -> List[Tuple[str, str]]:
+def google_search(self: Agent, query: str) -> List[Tuple[str, str]]:
     """
 
     A tool to search google with the provided query, and return a list of relevant summaries and URLs.
@@ -55,6 +33,18 @@ def google_search(self, query: str) -> List[Tuple[str, str]]:
         ]
     """
 
+    # imports must be inside the function
+    import os
+    import time
+    from concurrent.futures import ThreadPoolExecutor
+
+    import serpapi
+    from openai import OpenAI
+
+    from memgpt.credentials import MemGPTCredentials
+    from memgpt.data_sources.connectors import WebConnector
+    from memgpt.utils import printd
+
     printd("Starting google search:", query)
 
     def summarize_text(document_text: str, question: str) -> str:
@@ -100,9 +90,11 @@ def google_search(self, query: str) -> List[Tuple[str, str]]:
             links.append(data["link"])
         links = links[:5]
     except Exception as e:
-        printd(f"An error occurred with retrieving results: {e}")
+        print(f"An error occurred with retrieving results: {e}")
         return []
 
+    print("links", links)
+
     # retrieve text data from links
 
     def read_and_summarize_link(link):
@@ -112,7 +104,7 @@ def google_search(self, query: str) -> List[Tuple[str, str]]:
             printd(f"Time taken to retrieve text data: {time.time() - st}")
             # summarize text data
             st = time.time()
-            summary = summarize_text(document_text, query)
+            summary = summarize_text(document_text[: 16000 - 500], query)
             printd(f"Time taken to summarize text data: {time.time() - st}, length: {len(document_text)}")
             printd(link)
             if summary is not None:
@@ -127,7 +119,7 @@ def google_search(self, query: str) -> List[Tuple[str, str]]:
                 future = executor.submit(read_and_summarize_link, link)
                 futures.append(future)
         response = [future.result() for future in futures if future.result() is not None]
-        printd(f"Time taken: {time.time() - st}")
+        print(f"Time taken: {time.time() - st}")
         # response = []
         # connector = WebConnector(links)
         # for document_text, document_metadata in connector.generate_documents():
@@ -135,7 +127,58 @@ def google_search(self, query: str) -> List[Tuple[str, str]]:
         #    summary = summarize_text(document_text, query)
         #    if summary is not None:
         #        response.append((summary,  document_metadata["url"]))
+        print("Response:", response)
         return response
     except Exception as e:
-        printd(f"An error occurred with retrieving text data: {e}")
+        print(f"An error occurred with retrieving text data: {e}")
         return []
+
+
+def main():
+
+    # Create a `LocalClient` (you can also use a `RESTClient`, see the memgpt_rest_client.py example)
+    client = create_client()
+
+    # create tool
+    search_tool = client.create_tool(google_search, name="google_search")
+    print(f"Created tool: {search_tool.name} with ID {str(search_tool.id)}")
+    print(f"Tool schema: {json.dumps(search_tool.json_schema, indent=4)}")
+
+    # google search persona
+    persona = """
+
+    My name is MemGPT.
+
+    I am a personal assistant who answers a user's questionien using google web searches. When a user asks me a question and the answer is not in my context, I will use a tool called google_search which will search the web and return relevant summaries and the link they correspond to. It is my job to construct the best query to input into google_search based on the user's question, and to aggregate the response of google_search construct a final answer that also references the original links the information was pulled from. Here is an example:
+
+    ---
+
+    User: Who founded OpenAI?
+    MemGPT: OpenAI was founded by  Ilya Sutskever, Greg Brockman, Trevor Blackwell, Vicki Cheung, Andrej Karpathy, Durk Kingma, Jessica Livingston, John Schulman, Pamela Vagata, and Wojciech Zaremba, with Sam Altman and Elon Musk serving as the initial Board of Directors members. [1][2]
+
+    [1] https://www.britannica.com/topic/OpenAI
+    [2] https://en.wikipedia.org/wiki/OpenAI
+
+    ---
+
+    Don’t forget - inner monologue / inner thoughts should always be different than the contents of send_message! send_message is how you communicate with the user, whereas inner thoughts are your own personal inner thoughts.
+    """
+
+    # Create an agent
+    agent_state = client.create_agent(
+        name="my_agent3", memory=ChatMemory(human="My name is Sarah.", persona=persona), tools=[search_tool.name]
+    )
+    print(f"Created agent: {agent_state.name} with ID {str(agent_state.id)}")
+
+    # Send a message to the agent
+    print(f"Created agent: {agent_state.name} with ID {str(agent_state.id)}")
+    send_message_response = client.user_message(agent_id=agent_state.id, message="What is the weather in Berkeley?")
+    print(f"Recieved response: \n{json.dumps(send_message_response.messages, indent=4)}")
+
+    # Delete agent
+    client.delete_agent(agent_id=agent_state.id)
+    print(f"Deleted agent: {agent_state.name} with ID {str(agent_state.id)}")
+
+
+if __name__ == "__main__":
+    main()