agent.py

from langchain.agents import initialize_agent, AgentType
from langchain.llms import HuggingFaceHub
from langchain.tools import Tool
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
import os

class CrawlerAgent:
    def __init__(self, goal):
        self.llm = HuggingFaceHub(repo_id="meta-llama/Llama-3.1-8B", huggingfacehub_api_token= os.getenv())
        self.goal = goal
        self.tools = [
            Tool(name="Scrape", func=self._scrape, description="Scrape URL to markdown and extract links"),
            Tool(name="Decide", func=self._decide_action, description="Choose next action based on state")
        ]
        self.agent = initialize_agent(self.tools, self.llm, agent=AgentType.REACT_DESCRIPTION, verbose=True)

    def plan_sub_urls(self, start_url, num):
        prompt = f"Plan {num} parallel sub-paths from {start_url} for goal: {self.goal}. Return as comma-separated URLs."
        return self.agent.run(prompt).split(', ')

    def _scrape(self, url):
        from firecrawl import FirecrawlApp
        app = FirecrawlApp(api_key= os.getenv())
        data = app.scrape_url(url)
        return {'markdown': data['markdown'], 'links': self._extract_links(data['markdown'])}

    def _decide_action(self, state):
        prompt = f"State: {state}. Goal: {self.goal}. Choose action (0 to {len(state['links'])-1}) or -1 to stop."
        return int(self.agent.run(prompt).split()[-1])  # Parse last number

    def _extract_links(self, markdown):
        return [line.split('](')[1][:-1] for line in markdown.split('\n') if '](' in line]

    def crawl(self, env):
        state = env.reset()
        results = []
        while True:
            data = env._scrape(env.current_url)
            action = self._decide_action(data)
            if action < 0 or action >= len(data['links']):
                break
            state, reward, done, info = env.step(action, data['links'])
            results.append({"reward": reward, "markdown": info['markdown']})
            if done:
                break
        return results