geekphoenix's picture
Update crew.py
bf5695d verified
# # crew.py file of design_to_code
from crewai import Agent, Crew, Process, Task, LLM
from crewai.project import CrewBase, agent, crew, task
from tools.custom_tool import landing_ai_document_analysis
from dotenv import load_dotenv
import os
# Load environment variables
load_dotenv()
@CrewBase
class DocProcessing():
"""DocProcessing crew"""
agents_config = 'config/agents.yaml'
tasks_config = 'config/tasks.yaml'
def __init__(self, anthropic_api_key: str = None):
"""
Initialize DocProcessing crew with API key
Args:
anthropic_api_key: The Anthropic API key to use for LLM
"""
# Use provided API key or fall back to environment variable
api_key = anthropic_api_key or os.getenv("ANTHROPIC_API_KEY")
if not api_key:
raise ValueError("Anthropic API key is required. Please provide it as parameter or set ANTHROPIC_API_KEY environment variable.")
self.llm = LLM(
model="claude-3-haiku-20240307",
api_key=api_key,
temperature=0,
)
@agent
def document_analyst(self) -> Agent:
return Agent(
config=self.agents_config['document_analyst'],
tools=[landing_ai_document_analysis],
llm=self.llm,
)
@agent
def developer_agent(self) -> Agent:
return Agent(
config=self.agents_config['developer_agent'],
llm=self.llm,
)
@task
def document_analysis(self) -> Task:
return Task(
config=self.tasks_config['document_analysis'],
)
@task
def code_implementation(self) -> Task:
return Task(
config=self.tasks_config['code_implementation'],
)
@crew
def crew(self) -> Crew:
"""Creates the DocProcessing crew"""
return Crew(
agents=self.agents, # Automatically created by the @agent decorator
tasks=self.tasks, # Automatically created by the @task decorator
process=Process.sequential,
verbose=True,
)