Spaces:
Sleeping
Sleeping
File size: 3,094 Bytes
4327cda |
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 |
import os
from typing import Optional, Tuple, List, Union
from PIL import Image
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from .image_dryer import ImageDryer
class DryingAgent:
def __init__(self):
"""Initialize the DryingAgent with chat model and image processor."""
self.chat_model = ChatOpenAI(
base_url="https://openrouter.ai/api/v1",
model_name="google/gemini-pro",
openai_api_key=os.getenv("OPENROUTER_API_KEY"),
temperature=0.7
)
self.image_dryer = ImageDryer()
self.chat_history: List[Union[HumanMessage, AIMessage]] = []
self.current_image: Optional[Image.Image] = None
self.processed_image: Optional[Image.Image] = None
# System prompt for the agent
self.system_prompt = SystemMessage(content="""You are a helpful assistant specialized in drying items.
Your main task is to help users dry various items and provide advice about drying processes.
When users provide images, you should analyze them and suggest appropriate drying methods.
Always maintain a professional and helpful tone while focusing on drying-related queries.""")
def process_message(self, message: str, image: Optional[Image.Image] = None) -> Tuple[list, Optional[Image.Image]]:
"""Process a user message and optional image, return response and processed image."""
try:
if not message or not isinstance(message, str):
raise ValueError("Message must be a non-empty string")
self.current_image = image
messages = [self.system_prompt] + self.chat_history + [HumanMessage(content=message)]
response = self.chat_model.invoke(messages)
response_content = response.content if hasattr(response, 'content') else str(response)
if image is not None:
self.processed_image = self.image_dryer.process_image(image)
else:
self.processed_image = None
self.chat_history.append(HumanMessage(content=message))
self.chat_history.append(AIMessage(content=response_content))
if len(self.chat_history) > 20:
self.chat_history = self.chat_history[-20:]
return [
{"role": "user", "content": message},
{"role": "assistant", "content": response_content}
], self.processed_image
except ValueError as ve:
return [{"role": "assistant", "content": f"Invalid input: {str(ve)}"}], None
except Exception as e:
print(f"Error in process_message: {str(e)}")
return [{"role": "assistant", "content": f"An error occurred: {str(e)}"}], None
def reset(self):
"""Reset the agent's state."""
self.chat_history = []
self.current_image = None
self.processed_image = None |