120 lines
3.2 KiB
Python
120 lines
3.2 KiB
Python
from llama_cpp import Llama
|
|
from langchain_core.language_models import LLM
|
|
from langchain.prompts import PromptTemplate
|
|
from langchain.agents import AgentExecutor, create_react_agent
|
|
from langchain.tools import Tool
|
|
from typing import Optional, List
|
|
from googlesearch import search
|
|
from bs4 import BeautifulSoup
|
|
import requests
|
|
# --- LLaMA-CPP model setup ---
|
|
llm_model = Llama(
|
|
model_path="ReAction-1.5B.Q5_K_M.gguf",
|
|
n_ctx=2**21,
|
|
n_threads=8,
|
|
use_mlock=True,
|
|
verbose=False
|
|
)
|
|
|
|
# --- LangChain wrapper ---
|
|
class LlamaCppLLM(LLM):
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
return "llama-cpp"
|
|
|
|
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
|
|
result = llm_model(
|
|
prompt,
|
|
stop=stop,
|
|
max_tokens=1024,
|
|
echo=False,
|
|
|
|
)
|
|
output = result["choices"][0]["text"].strip()
|
|
return output
|
|
|
|
custom_llm = LlamaCppLLM()
|
|
|
|
# --- Tool definition ---
|
|
def extract_text(html):
|
|
soup = BeautifulSoup(html, "html.parser")
|
|
# Remove scripts and styles
|
|
for tag in soup(["script", "style"]):
|
|
tag.decompose()
|
|
text = soup.get_text(separator=" ", strip=True)
|
|
return text
|
|
def weather(city: str) -> str:
|
|
return f"The weather in {city} is Sunny."
|
|
def search_google(query:str):
|
|
print(f"serach tool called {query}")
|
|
return [url for url in search(query,stop=5)]
|
|
def get_webhtml(url):
|
|
print(url)
|
|
url = url.replace("'","")
|
|
if url:
|
|
r = requests.get(url)
|
|
else:
|
|
return None
|
|
return extract_text(r.text) if r.status_code == 200 else None
|
|
weather_tool = Tool(
|
|
name="weather",
|
|
func=weather,
|
|
description="Use this tool to get the weather of a given city. Input should be the city name."
|
|
)
|
|
search_tool = Tool(
|
|
name = "search_tool",
|
|
func=search_google,
|
|
description="Returns the first 5 results from google with a specific query"
|
|
)
|
|
html_tool = Tool(
|
|
name="get_webhtml",
|
|
func=get_webhtml,
|
|
description="Given a url, it provides the HTML for it. If it returns None, the website isn't available."
|
|
)
|
|
|
|
# --- Custom ReAct prompt ---
|
|
prompt_template = PromptTemplate.from_template("""
|
|
You are a helpful assistant that can use tools to answer questions.
|
|
|
|
TOOLS:
|
|
{tools}
|
|
|
|
FORMAT:
|
|
Question: the input question
|
|
Thought: think about what to do
|
|
Action: pick one of [{tool_names}]
|
|
Action Input: "<the input to the action>"
|
|
|
|
Observation: result of the action
|
|
... (you can repeat Thought/Action/Observation until you reach a Final Answer) ...
|
|
... (Also, only use Action Input AFTER you use Action, so before you use Action input again, you have to call a tool again. Reach a final Answer before 3 iterations) ...
|
|
Thought: I now know the final answer
|
|
Final Answer: <the final answer>
|
|
|
|
|
|
NEW QUESTION:
|
|
|
|
Question: {input}
|
|
{agent_scratchpad}
|
|
""")
|
|
|
|
# --- Create ReAct agent ---
|
|
react_agent = create_react_agent(
|
|
llm=custom_llm,
|
|
tools=[search_tool,html_tool],
|
|
prompt=prompt_template
|
|
)
|
|
|
|
agent = AgentExecutor(
|
|
agent=react_agent,
|
|
tools=[search_tool,html_tool],
|
|
verbose=True,
|
|
handle_parsing_errors=True,
|
|
max_iterations=5
|
|
)
|
|
|
|
# --- Run it ---
|
|
response = agent.invoke({"input": "What is the weather in San Francisco"})
|
|
print("\nFinal Output:\n", response)
|
|
|