From cdfe463830fed9b55ce6d5cdeaf2226d01f77161 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fl=C3=A1vio=20Juvenal?= Date: Fri, 27 Sep 2024 18:37:23 -0300 Subject: [PATCH] Incomplete --- django_ai_assistant/helpers/assistants.py | 48 ++++++++++---------- example/demo/views.py | 1 - example/example/settings.py | 3 ++ example/movies/ai_assistants.py | 53 +++++++++++++++-------- 4 files changed, 60 insertions(+), 45 deletions(-) diff --git a/django_ai_assistant/helpers/assistants.py b/django_ai_assistant/helpers/assistants.py index dd85204..28da94c 100644 --- a/django_ai_assistant/helpers/assistants.py +++ b/django_ai_assistant/helpers/assistants.py @@ -75,6 +75,8 @@ class AIAssistant(abc.ABC): # noqa: F821 """ temperature: float = 1.0 """Temperature to use for the assistant LLM model.\nDefaults to `1.0`.""" + tool_max_concurrency: int = 1 + """Maximum number of tools to run concurrently / in parallel.\nDefaults to `1` (no concurrency).""" has_rag: bool = False """Whether the assistant uses RAG (Retrieval-Augmented Generation) or not.\n Defaults to `False`. @@ -430,7 +432,7 @@ def as_graph(self, thread_id: Any | None = None) -> Runnable[dict, dict]: llm_with_tools = llm.bind_tools(tools) if tools else llm def custom_add_messages(left: list[BaseMessage], right: list[BaseMessage]): - result = add_messages(left, right) + result = add_messages(left, right) # type: ignore if message_history: messages_to_store = [ @@ -447,40 +449,30 @@ class AgentState(TypedDict): messages: Annotated[list[AnyMessage], custom_add_messages] input: str # noqa: A003 context: str - output: str + output: Any def setup(state: AgentState): - messages: list[AnyMessage] = [SystemMessage(content=self.get_instructions())] + system_prompt = self.get_instructions() if self.structured_output: - schema = None - # If Pydantic if inspect.isclass(self.structured_output) and issubclass( self.structured_output, BaseModel ): schema = json.dumps(self.structured_output.model_json_schema()) - - schema_information = ( - f"JSON will have the following schema:\n\n{schema}\n\n" if schema else "" - ) - tools_information = "Gather information using tools. " if tools else "" - - # The assistant won't have access to the schema of the structured output before - # the last step of the chat. This message gives visibility about what fields the - # response should have so it can gather the necessary information by using tools. - messages.append( - HumanMessage( - content=( - "In the last step of this chat you will be asked to respond in JSON. " - + schema_information - + tools_information - + "Don't generate JSON until you are explicitly told to. " - ) + schema_information = ( + f"Your JSON output must have the following schema:\n{schema}\n" + if schema + else "" ) - ) + json_info = ( + "In the last step of this chat you will be asked to respond in JSON. " + + schema_information + + "Don't generate JSON until you are explicitly told to. " + ) + system_prompt += "\n" + json_info - return {"messages": messages} + return {"messages": [SystemMessage(content=system_prompt)]} def history(state: AgentState): history = message_history.messages if message_history else [] @@ -523,12 +515,14 @@ def tool_selector(state: AgentState): def record_response(state: AgentState): if self.structured_output: + # Structured output must happen in the end, to avoid disabling tool calling. + # Tool calling + structured output is not supported by OpenAI: llm_with_structured_output = self.get_structured_output_llm() response = llm_with_structured_output.invoke( [ *state["messages"], HumanMessage( - content="Use the information gathered in the conversation to answer." + content="Use the information gathered in the conversation to answer with JSON." ), ] ) @@ -581,7 +575,9 @@ def invoke(self, *args: Any, thread_id: Any | None, **kwargs: Any) -> dict: structured like `{"output": "assistant response", "history": ...}`. """ graph = self.as_graph(thread_id) - return graph.invoke(*args, **kwargs) + config = kwargs.pop("config", {}) + config["max_concurrency"] = config.pop("max_concurrency", self.tool_max_concurrency) + return graph.invoke(*args, config=config, **kwargs) @with_cast_id def run(self, message: str, thread_id: Any | None = None, **kwargs: Any) -> str: diff --git a/example/demo/views.py b/example/demo/views.py index d2b91b6..c52ef88 100644 --- a/example/demo/views.py +++ b/example/demo/views.py @@ -119,5 +119,4 @@ def get(self, request, *args, **kwargs): a = TourGuideAIAssistant() data = a.run(f"My coordinates are: ({coordinates})") - return JsonResponse(data.model_dump()) diff --git a/example/example/settings.py b/example/example/settings.py index a16b829..7a5fac0 100644 --- a/example/example/settings.py +++ b/example/example/settings.py @@ -180,4 +180,7 @@ # Example specific settings: WEATHER_API_KEY = os.getenv("WEATHER_API_KEY") # get for free at https://www.weatherapi.com/ +BRAVE_SEARCH_API_KEY = os.getenv( + "BRAVE_SEARCH_API_KEY" +) # get for free at https://brave.com/search/api/ DJANGO_DOCS_BRANCH = "stable/5.0.x" diff --git a/example/movies/ai_assistants.py b/example/movies/ai_assistants.py index 8db5552..d4f25b0 100644 --- a/example/movies/ai_assistants.py +++ b/example/movies/ai_assistants.py @@ -1,32 +1,42 @@ from typing import Sequence +from django.conf import settings from django.db import transaction from django.db.models import Max from django.utils import timezone import requests -from langchain_community.tools.tavily_search import TavilySearchResults +from langchain_community.tools import BraveSearch from langchain_core.tools import BaseTool +from pydantic import BaseModel from django_ai_assistant import AIAssistant, method_tool from movies.models import MovieBacklogItem +class IMDbMovie(BaseModel): + imdb_url: str + imdb_rating: float + scrapped_imdb_page_markdown: str + + # Note this assistant is not registered, but we'll use it as a tool on the other. # This one shouldn't be used directly, as it does web searches and scraping. -class IMDbURLFinderTool(AIAssistant): - id = "imdb_url_finder" # noqa: A003 +class IMDbScraper(AIAssistant): + id = "imdb_scraper" # noqa: A003 instructions = ( - "You're a tool to find the IMDb URL of a given movie. " - "Use the Tavily Search API to find the IMDb URL. " + "You're a tool to find the IMDb URL of a given movie, " + "and scrape this URL to get the movie rating and other information.\n" + "Use the search tool to find the IMDb URL. " "Make search queries like: \n" "- IMDb page of The Matrix\n" "- IMDb page of The Godfather\n" "- IMDb page of The Shawshank Redemption\n" - "Then check results and provide only the IMDb URL to the user." + "Then check results, scape the IMDb URL, process the page, and produce a JSON output." ) - name = "IMDb URL Finder" - model = "gpt-4o-mini" + name = "IMDb Scraper" + model = "gpt-4o" + structured_output = IMDbMovie def get_instructions(self): # Warning: this will use the server's timezone @@ -35,9 +45,16 @@ def get_instructions(self): current_date_str = timezone.now().date().isoformat() return f"{self.instructions} Today is: {current_date_str}." + @method_tool + def scrape_imdb_url(self, url: str) -> str: + """Scrape the IMDb URL and return the content as markdown.""" + return requests.get("https://r.jina.ai/" + url, timeout=20).text[:10000] + def get_tools(self) -> Sequence[BaseTool]: return [ - TavilySearchResults(), + BraveSearch.from_api_key( + api_key=settings.BRAVE_SEARCH_API_KEY, search_kwargs={"count": 5} + ), *super().get_tools(), ] @@ -47,6 +64,11 @@ class MovieRecommendationAIAssistant(AIAssistant): instructions = ( "You're a helpful movie recommendation assistant. " "Help the user find movies to watch and manage their movie backlogs. " + "Use the provided tools for that.\n" + "Note the backlog is stored in a DB. " + "When managing the backlog, you must call the tools, to keep the sync with the DB. " + "The backlog has an order, and you should respect it. Call `reorder_backlog` when necessary.\n" + "Include the IMDb URL and rating of the movies when displaying the backlog.\n" "Ask the user if they want to add your recommended movies to their backlog, " "but only if the movie is not on the user's backlog yet." ) @@ -70,18 +92,13 @@ def get_instructions(self): def get_tools(self) -> Sequence[BaseTool]: return [ - TavilySearchResults(), - IMDbURLFinderTool().as_tool(description="Tool to find the IMDb URL of a given movie."), + BraveSearch.from_api_key( + api_key=settings.BRAVE_SEARCH_API_KEY, search_kwargs={"count": 5} + ), + IMDbScraper().as_tool(description="Tool to get the IMDb data a given movie."), *super().get_tools(), ] - @method_tool - def scrape_imdb_url(self, url: str) -> str: - """Scrape the IMDb URL and return the content as markdown. - Use this to get more info about a movie / show, including its rating, plot, cast, etc.""" - - return requests.get("https://r.jina.ai/" + url, timeout=20).text[:10000] - @method_tool def get_movies_backlog(self) -> str: """Get what movies are on user's backlog."""