Spaces:
Sleeping
Sleeping
File size: 3,789 Bytes
2fcb36b |
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 |
import base64
import httpx
from pymongo import MongoClient
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage
import os
import re
import json
from dotenv import load_dotenv
load_dotenv()
MONGO_URI = os.getenv("MONGO_URI")
DB_NAME = os.getenv("DB_NAME")
COLLECTION_NAME = os.getenv("COLLECTION_NAME")
FLASH_API = os.getenv("FLASH_API")
mongo_client = MongoClient(MONGO_URI)
db = mongo_client[DB_NAME]
collection = db[COLLECTION_NAME]
collection2=db['about_company']
model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0, max_tokens=None, google_api_key=FLASH_API)
about_company_doc=collection2.find_one({"type":"about_company"})
if about_company_doc:
about_company=about_company_doc.get('company_description','')
system_prompt_text = f"""Given is an image related to a company. Your task is to analyze the image, identify any text or notable visual elements, and provide a comprehensive, direct description of the image's contents, focusing on what it represents without abstract language or additional commentary. The response must be concise and focused, using only descriptive nouns and adjectives. If the image cannot be clearly described, respond with 'None.'
Company information is given below to understand the context.
- About Company: {about_company}
Ensure you give valid JSON structure.
Expected Output format : {{"description":"String"}}
"""
def process_image_using_llm(image_url):
try:
# Download and encode the image
image_data = base64.b64encode(httpx.get(image_url).content).decode("utf-8")
# Create the message with a system prompt and image
message = HumanMessage(
content=[
{"type": "text", "text": system_prompt_text},
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}},
],
)
# Invoke the model to extract information from the image
response = model.invoke([message])
print(f"llm response for {image_url} is : {response}")
# Use regex to extract JSON part of the response
match = re.search(r"\{.*\}", response.content.strip())
if match:
json_data = match.group(0) # Extract JSON-like content as a string
data = json.loads(json_data) # Load as JSON
# Get the description from the JSON data
description = data.get("description", "None").strip()
# Check if the description is "None"
if description == "None":
# Update MongoDB with can_find_description as False
collection.update_one(
{"object_url": image_url},
{"$set": {"can_find_description": False}}
)
print(f"Marked {image_url} as can_find_description: False")
return False
else:
# Update MongoDB with the description and set can_find_description to True
collection.update_one(
{"object_url": image_url},
{"$set": {"description": description, "can_find_description": True}}
)
print("Saved description to MongoDB")
return True
else:
print(f"No valid JSON found in the response for {image_url}. Marking as can_find_description: False")
collection.update_one(
{"object_url": image_url},
{"$set": {"can_find_description": False}}
)
return False
except Exception as e:
print(f"Error processing {image_url}: {e}")
return False
|