Merge branch 'main' into distributed

examples/llm_lottery.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+
+# Tip: Most fun when run with `-v` or `-vv` flag!
+#
+# Small educational example on how to use tools.
+# The example provides LLM with two tools:
+# * One to buy a lottery ticket.
+# * Another to check a lottery ticket.
+# 
+# Then we run LLM to enjoy watching how it plays the lottery.
+#
+# Additionally, it provides:
+# `loop_llm_and_tools` - an example of a handy boilerplate function
+#                        that helps to run tools until the task is completed.
+# `main` with argparse and -v|--verbose flag -
+#        An example of how to make simple code convenient to switch between
+#        different verbosity levels from the command line, e.g. for this script:
+#        * `-v`  prints progress information to stderr.
+#        * `-vv` also enables verbosity for ell.init.
+
+import ell
+from ell import Message
+import argparse
+import sys
+from typing import List
+from pydantic import Field
+
+VERBOSE = False
+WINNING_NUMBER = 12  # Hardcoded winning number
+
+@ell.tool(strict=True)
+def buy_lottery_ticket(number: int = Field(description="Number to play in the lottery (0-16).")):
+    """Buy a lottery ticket with the given number."""
+    if VERBOSE:
+        sys.stderr.write(f"Calling tool: buy_lottery_ticket({number})\n")
+    return f"Bought lottery ticket with number {number}"
+
+@ell.tool(strict=True)
+def check_lottery_result(number: int = Field(description="Number to check against the lottery result.")):
+    """Check if the given number wins the lottery."""
+    if VERBOSE:
+        sys.stderr.write(f"Calling tool: check_lottery_result({number})\n")
+    if number == WINNING_NUMBER:
+        return "Congratulations! You won the lottery!"
+    else:
+        return "Sorry, your number did not win. Try again!"
+
+@ell.complex(model="gpt-4o-mini", tools=[buy_lottery_ticket, check_lottery_result])
+def play_lottery(message_history: List[Message]) -> List[Message]:
+    if VERBOSE:
+        last_msg = message_history[-1].text
+        sys.stderr.write(f"Calling LMP: play_lottery('{last_msg}')\n")
+    return [
+        ell.system("You are an AI assistant that plays the lottery. Buy a lottery ticket with a number between 0 and 16, then check if it wins. If it doesn't win, try again with a different number.")
+    ] + message_history
+
+def loop_llm_and_tools(f, message_history, max_iterations=100):
+    iteration = 0
+    while iteration < max_iterations:
+        response_message = f(message_history)
+        message_history.append(response_message)
+
+        if response_message.tool_calls:
+            tool_call_response = response_message.call_tools_and_collect_as_message()
+            message_history.append(tool_call_response)
+
+            # Check if we've won the lottery
+            if "Congratulations" in tool_call_response.text:
+                break
+        else:
+            break
+        iteration += 1
+    return message_history
+
+def main():
+    parser = argparse.ArgumentParser(description='Play the lottery until winning.')
+    parser.add_argument('-v', '--verbose', action='count', default=0, help='Increase verbosity level')
+    args = parser.parse_args()
+
+    global VERBOSE
+    VERBOSE = args.verbose > 0
+
+    ell.init(verbose=(args.verbose > 1), store='./logdir', autocommit=True)
+
+    message_history = []
+    message_history.append(ell.user("Let's play the lottery until we win!"))
+
+    message_history = loop_llm_and_tools(play_lottery, message_history)
+
+    print("\nLottery Game Results:")
+    for message in message_history:
+        if message.role == "assistant" or message.role == "function":
+            print(f"{message.role.capitalize()}: {message.text}")
+
+    print(f"\nTotal attempts: {len(message_history) // 2 - 1}")
+
+if __name__ == "__main__":
+    main()
+

examples/wikipedia_mini_rag.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+
+# Educational Example: Using LLM with Wikipedia Tools
+#
+# This script demonstrates how to use an LLM agent with tools to interact with Wikipedia.
+# It provides two main functionalities:
+# * Searching Wikipedia for relevant pages.
+# * Fetching and reading the content of a Wikipedia page.
+#
+# The script uses `lynx --dump` to obtain a textual representation of web pages.
+#
+# Workflow:
+# 1. The AI agent searches Wikipedia based on user queries and suggests a page to read.
+#     * Agent keeps searching, trying different kewords, until sees some promising result.
+# 2. The agent uses tool to read the page content to answer the user's query.
+#
+# This example is designed to be educational and is suitable for inclusion in an `examples/` directory.
+# It illustrates the integration of LLMs with external tools to perform complex tasks.
+#
+# For those interested in understanding the inner workings, try running the script with the `-v` or `-vv` flags.
+# These flags will provide additional insights into the process and can be very helpful for learning purposes.
+#
+# Bonus Task: Consider looking at `llm_lottery.py` which uses `loop_llm_and_tools` to allow the LLM to call tools
+# as long as needed to accomplish a task. Think about how you might modify this script to search and read Wikipedia
+# pages until it knows enough to provide a sufficient response to a query, such as a comparison between multiple cities.
+# (spoiler, example solution: https://gist.github.com/gwpl/27715049d41ec829f21014f3b243850a )
+
+import argparse
+import subprocess
+import sys
+from functools import partial
+import urllib.parse
+
+import ell
+from ell import Message
+from pydantic import Field
+
+eprint = partial(print, file=sys.stderr)
+
+VERBOSE = False
+
+@ell.tool(strict=True)
+def search_wikipedia(keywords: str = Field(description="Kewords for wikipedia search engine.")):
+    """Search Wikipedia and return a list of search results and links."""
+    if VERBOSE:
+        eprint(f"Calling tool: search_wikipedia('{keywords}')")
+
+    encoded_query = urllib.parse.quote(keywords)
+    cmd = f"lynx --dump 'https://en.m.wikipedia.org/w/index.php?search={encoded_query}'"
+    result = subprocess.run(cmd, shell=True, capture_output=True)
+    return result.stdout.decode('ISO-8859-1')[:65536]
+
+@ell.tool(strict=True)
+def wikipedia_page_content(wiki_page_url: str):
+    """Fetch the content of a Wikipedia page given its URL."""
+    if VERBOSE:
+        eprint(f"Calling tool: wikipedia_page_content('{wiki_page_url}')")
+
+    cmd = f"lynx --dump '{wiki_page_url}'"
+    result = subprocess.run(cmd, shell=True, capture_output=True)
+    return result.stdout.decode('ISO-8859-1')[:65536]
+
+@ell.complex(model="gpt-4o-mini", tools=[search_wikipedia])
+def search_wikipedia_and_suggest_page_to_read(message_history: list[Message]) -> list[Message]:
+    if VERBOSE:
+        last_msg = message_history[-1].text
+        if len(last_msg) > 100:
+            last_msg = last_msg[:100] + "..."
+        eprint(f"Calling LMP: search_wikipedia_and_suggest_page_to_read('{last_msg}')")
+    return [
+        ell.system("You are an AI assistant that searches Wikipedia and suggests URL of wikipedia page to read. Return ONLY URL. Call tool that takes keywords for Wikipedia search engine, to get wikipedia search engine results. Based on results suggest one best, most promisting/relevent URL to read.")
+    ] + message_history
+
+@ell.complex(model="gpt-4o", tools=[wikipedia_page_content])
+def answer_query_by_reading_wikipedia_page(message_history: list[Message]) -> list[Message]:
+    if VERBOSE:
+        last_msg = message_history[-1].text
+        eprint(f"Calling LMP: answer_query_by_reading_wikipedia_page({last_msg})")
+    return [
+        ell.system("You are an AI assistant that reads a Wikipedia page and provides answers based on the information found. Read the content of the Wikipedia page at provided URL. Call tool to fetch content of page. Use the content to provide a comprehensive answer to user query. Include quotes on which answer is based on.")
+    ] + message_history
+
+def loop_llm_and_tools(f, message_history, max_iterations=10):
+    iteration = 0
+    while iteration < max_iterations:
+        response_message = f(message_history)
+        message_history.append(response_message)
+
+        if response_message.tool_calls:
+            tool_call_response = response_message.call_tools_and_collect_as_message()
+            message_history.append(tool_call_response)
+        else:
+            break
+        iteration += 1
+    return message_history
+
+def main():
+    parser = argparse.ArgumentParser(description='Search Wikipedia and answer questions.')
+    parser.add_argument('query', type=str, help='The query to search for on Wikipedia')
+    parser.add_argument('-v', '--verbose', action='count', default=0, help='Increase verbosity level')
+    args = parser.parse_args()
+
+    global VERBOSE
+    VERBOSE = args.verbose > 0
+
+    ell.init(verbose=(args.verbose > 1), store='./logdir', autocommit=True)
+
+    if args.verbose > 0:
+        eprint(f"Provided Query: {args.query}")
+
+    message_history = []
+    message_history.append(ell.user(args.query))
+
+    if args.verbose > 1:
+        eprint(f"message_history as early stage = {message_history}")
+
+    message_history = loop_llm_and_tools(search_wikipedia_and_suggest_page_to_read, message_history)
+
+    url = message_history[-1].text
+    eprint(f"Intermediate Result: {url}")
+
+    message_history = loop_llm_and_tools(answer_query_by_reading_wikipedia_page, message_history)
+
+    eprint("Result:")
+    print(f"{message_history[-1].text}\n")
+
+if __name__ == "__main__":
+    main()
+
+
+# This script does the following:
+#
+# 1. Defines two tools: `search_wikipedia` and `wikipedia_page_content` using the `@ell.tool()` decorator.
+#    These tools use the `lynx` command-line browser to fetch search results and page content from Wikipedia.
+#
+# 2. Implements two complex functions using the `@ell.complex` decorator:
+#    - `search_wikipedia_and_suggest_page_to_read`: Searches Wikipedia and suggests a page URL.
+#    - `answer_query_by_reading_wikipedia_page`: Reads a Wikipedia page and answers the user's query.
+#
+# 3. The `main` function sets up argument parsing, initializes ell with the appropriate verbosity, and manages the
+#    interaction loop using `loop_llm_and_tools` to process the query and fetch results.
+#
+# 4. The script supports two levels of verbosity:
+#    - With `-v`, it prints progress information to stderr.
+#    - With `-vv`, it also enables verbosity for ell.init.
+#
+# 5. Finally, it prints the intermediate URL result and the final answer based on the Wikipedia page content.
+#
+# To use this script, you would run it from the command line like this:
+#
+# ```
+# python3 wikipedia_mini_rag.py "Your query here" -v
+# ```
+#
+# Make sure you have the `lynx` command-line browser installed on your system for this script to work properly.

src/ell/util/closure.py

@@ -83,7 +83,7 @@ def lexical_closure(
     while hasattr(func, "__ell_func__"):
         func = func.__ell_func__
 
-    source = getsource(func, lstrip=True)
+    source = getsource(func, lstrip=True, force=True)
     already_closed.add(hash(func))
 
     globals_and_frees = _get_globals_and_frees(func)