File size: 4,539 Bytes
6e5459c
6520972
6e5459c
391d348
 
6e5459c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
391d348
6e5459c
391d348
6e5459c
 
391d348
6e5459c
 
 
 
 
 
 
 
 
 
391d348
6e5459c
 
 
 
 
 
 
391d348
6e5459c
391d348
 
6e5459c
 
 
 
391d348
6e5459c
 
 
 
 
 
 
 
 
391d348
 
6e5459c
 
 
 
 
 
 
 
 
391d348
6e5459c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6520972
 
6e5459c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6520972
 
e960060
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import os
import sys
import logging
from typing import Optional, Any
from typing_extensions import override

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

def setup_environment():
    """Setup environment variables and cache directories"""
    try:
        cache_dir = "/tmp/hf"
        os.environ["TRANSFORMERS_CACHE"] = cache_dir
        os.environ["HF_HOME"] = cache_dir
        
        # Create cache directory if it doesn't exist
        os.makedirs(cache_dir, exist_ok=True)
        logger.info(f"Cache directory set up at {cache_dir}")
        return True
    except Exception as e:
        logger.error(f"Failed to setup environment: {str(e)}")
        return False

def initialize_model():
    """Initialize the model and pipeline with error handling"""
    try:
        from transformers import pipeline
        logger.info("Loading model pipeline...")
        
        pipe = pipeline(
            "text-generation",
            model="bigscience/bloom-560m",
            tokenizer="bigscience/bloom-560m",
            device=-1,
            max_new_tokens=256,
            force_download=True,
            low_cpu_mem_usage=True,
        )
        logger.info("Model pipeline loaded successfully")
        return pipe
    except Exception as e:
        logger.error(f"Failed to initialize model: {str(e)}")
        return None

def setup_agent(pipe):
    """Setup LangChain agent with the model pipeline"""
    try:
        from langchain_huggingface import HuggingFacePipeline  # Updated import
        from langchain.agents import initialize_agent
        from langchain.agents.agent_types import AgentType
        from langchain.tools import BaseTool

        # Initialize LLM
        llm = HuggingFacePipeline(pipeline=pipe)

        # Define the tool with proper type annotations
        class MachineryReportTool(BaseTool):
            name: str = "machinery_report"  # Type annotation added
            description: str = (  # Type annotation added
                "Generates a report on mini construction equipment including "
                "specifications and market analysis."
            )

            @override
            def _run(self, query: str) -> str:
                return """
                Mini Construction Equipment Report:
                1. Basic Gas-Powered Unit: $3,700, 14HP
                2. DRT 450: $5,500-6,300, Honda motor
                3. Mini Skid Steer: $15,000, Kubota diesel
                Market Analysis: 30-40% cheaper than US equivalents
                """

            @override
            def _arun(self, query: str) -> Any:
                raise NotImplementedError("Async not supported.")

        # Initialize agent with minimal settings
        agent = initialize_agent(
            tools=[MachineryReportTool()],
            llm=llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=True,
            max_iterations=1,
            early_stopping_method="generate"
        )
        logger.info("Agent initialized successfully")
        return agent
    except Exception as e:
        logger.error(f"Failed to setup agent: {str(e)}")
        return None

def run_agent_query(agent, query: str) -> Optional[str]:
    """Run a query through the agent with error handling"""
    try:
        logger.info("Running agent query...")
        result = agent.run(query)
        logger.info("Query completed successfully")
        return result
    except Exception as e:
        logger.error(f"Error running agent query: {str(e)}")
        return f"Error processing query: {str(e)}"

def main():
    """Main function to run the agent"""
    logger.info("Starting application...")
    
    # Setup environment
    if not setup_environment():
        return
    
    # Initialize model
    pipe = initialize_model()
    if pipe is None:
        return
    
    # Setup agent
    agent = setup_agent(pipe)
    if agent is None:
        return
    
    # Define query
    query = (
        "Please provide a brief overview of the mini construction equipment "
        "project, focusing on key specifications and market positioning."
    )
    
    # Run query
    result = run_agent_query(agent, query)
    
    # Print result
    if result:
        print("\n===== AGENT OUTPUT =====")
        print(result)
        print("=======================")

if __name__ == "__main__":
    main()