You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
74 lines
3.5 KiB
74 lines
3.5 KiB
import os |
|
from langchain.tools.base import BaseTool |
|
from langchain.callbacks.manager import CallbackManagerForToolRun |
|
import requests |
|
from typing import Optional, Dict, List |
|
from langchain_community.tools import DuckDuckGoSearchRun |
|
import random |
|
import time |
|
|
|
|
|
class MediaInfoSearchTool(BaseTool): |
|
name: str = "duckduckgo_search" |
|
description: str = '''Useful for searching the web using DuckDuckGo for information about \ |
|
movies and TV shows, actors and directors.''' |
|
|
|
# Class variable to track previous queries and sites |
|
movie_sites: List[str] = ["imdb.com", "rottentomatoes.com", "metacritic.com", "themoviedb.org", "filmaffinity.com"] |
|
|
|
def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None) -> str: |
|
"""Perform a DuckDuckGo search with site rotation queries.""" |
|
try: |
|
search_tool = DuckDuckGoSearchRun() |
|
result = "" |
|
# Randomly select 3 sites from the movie_sites list |
|
selected_sites = random.sample(self.movie_sites, 2) |
|
|
|
for movie_site in selected_sites: |
|
result += f"Searching for '{query}' on {movie_site}...\n" |
|
try: |
|
# Perform the search using DuckDuckGo |
|
result += search_tool.run(f"{query} site:{movie_site}") |
|
except Exception as e: |
|
result += f"Error searching on {movie_site}: {str(e)}\n" |
|
time.sleep(1) # Sleep for 1 second to avoid hitting the API too fast |
|
# Perform the search using DuckDuckGo |
|
result += search_tool.run(f"{query} site:{movie_site}") |
|
result += "\n\n" |
|
print(f"Searching for '{query}' on {movie_site}...\n") |
|
return result |
|
|
|
except Exception as e: |
|
return f"Error searching for '{query}': {str(e)}" |
|
|
|
class MoviesAdviceSearchTool(BaseTool): |
|
name: str = "movies_advice_search" |
|
description: str = '''Useful for searching the web using DuckDuckGo for movie recommendations and similar content to a given title or plot. |
|
prefer searching on trustworthy sites. |
|
Input should be a search query, and the tool will return relevant results.''' |
|
|
|
# Class variable to track recommendation sites |
|
recommendation_sites: List[str] = ["reddit.com/r/moviesuggestions", "tastedive.com", "letterboxd.com", "movielens.org", "flickmetrix.com", "justwatch.com"] |
|
|
|
def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None) -> str: |
|
"""Perform a DuckDuckGo search with site rotation queries.""" |
|
try: |
|
search_tool = DuckDuckGoSearchRun() |
|
result = "" |
|
# Randomly select 2 sites from the recommendation_sites list |
|
selected_sites = random.sample(self.recommendation_sites, 2) |
|
|
|
for rec_site in selected_sites: |
|
result += f"Searching for '{query}' on {rec_site}...\n" |
|
try: |
|
# Perform the search using DuckDuckGo |
|
result += search_tool.run(f"{query} site:{rec_site}") |
|
except Exception as e: |
|
result += f"Error searching on {rec_site}: {str(e)}\n" |
|
time.sleep(5) # Sleep for 1 second to avoid hitting the API too fast |
|
result += "\n\n" |
|
print(f"Searching for '{query}' on {rec_site}...\n") |
|
return result |
|
|
|
except Exception as e: |
|
return f"Error searching for '{query}': {str(e)}" |