Spaces:
Sleeping
Sleeping
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 a 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 | |