File size: 3,101 Bytes
458c8e2 |
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 |
import os
from typing import List, Optional
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv
load_dotenv()
def simple_chat(
question: str,
context: str,
image_paths: Optional[List[str]] = None
) -> str:
"""
Simple chat function that answers questions based on context and optional images.
Args:
question: User's question
context: Context information (e.g., dataset summary, analysis results)
image_paths: Optional list of image file paths to include
Returns:
AI response as a string
"""
try:
# Initialize LLM
llm = ChatGoogleGenerativeAI(
model="gemini-2.0-flash-exp",
temperature=0,
api_key=os.getenv("GOOGLE_API_KEY"),
)
# Build the prompt
prompt = f"""You are a helpful data analysis assistant.
Context:
{context}
User Question: {question}
Please provide a clear, concise answer based on the context provided."""
# Handle images if provided
if image_paths:
content = [{"type": "text", "text": prompt}]
for img_path in image_paths:
if os.path.exists(img_path):
import base64
with open(img_path, "rb") as f:
img_data = base64.b64encode(f.read()).decode()
# Determine image type
ext = os.path.splitext(img_path)[1].lower()
mime_type = {
'.png': 'image/png',
'.jpg': 'image/jpeg',
'.jpeg': 'image/jpeg',
'.gif': 'image/gif',
'.webp': 'image/webp'
}.get(ext, 'image/png')
content.append({
"type": "image_url",
"image_url": f"data:{mime_type};base64,{img_data}"
})
message = HumanMessage(content=content)
else:
message = HumanMessage(content=prompt)
# Get response
response = llm.invoke([message])
# Extract text from response
if hasattr(response, 'content'):
return str(response.content)
return str(response)
except Exception as e:
return f"Error: {str(e)}"
# Example usage:
if __name__ == "__main__":
# Simple text-only example
context = """
Dataset: Customer Sales Data
- 1000 rows, 15 columns
- Label: purchase_made (binary)
- Task: Classification
- Missing values: 5% in age column
"""
question = "What's the main task for this dataset?"
response = simple_chat(question, context)
print(response)
# With images
question2 = "What do you see in the visualization?"
response2 = simple_chat(question2, context, image_paths=["/path/to/plot.png"])
print(response2) |