Fixed metadata. Fixed image retention.

This commit is contained in:
Andrew Tait Gehrhardt 2024-06-18 21:05:18 -04:00
parent cc936574e5
commit 7d817b6b6a

View File

@ -1,10 +1,10 @@
""" """
title: HomeAssistant Filter Pipeline title: Ollama Dynamic Vision Pipeline
author: Andrew Tait Gehrhardt author: Andrew Tait Gehrhardt
date: 2024-06-15 date: 2024-06-18
version: 1.0 version: 1.0
license: MIT license: MIT
description: A pipeline for controlling Home Assistant entities based on their easy names. Only supports lights at the moment. description: A pipeline for dynamically processing images when current model is a text only model
requirements: pydantic, aiohttp requirements: pydantic, aiohttp
""" """
@ -12,15 +12,15 @@ from typing import List, Optional
from pydantic import BaseModel from pydantic import BaseModel
import json import json
import aiohttp import aiohttp
from utils.pipelines.main import get_last_user_message, get_last_assistant_message from utils.pipelines.main import get_last_user_message
class Pipeline: class Pipeline:
class Valves(BaseModel): class Valves(BaseModel):
pipelines: List[str] = [] pipelines: List[str] = []
priority: int = 0 priority: int = 0
target_user_roles: List[str] = ["admin", "user"]
vision_model: str = "llava" vision_model: str = "llava"
ollama_base_url: str = "http://host.docker.interal:11434" ollama_base_url: str = ""
model_to_override: str = ""
def __init__(self): def __init__(self):
self.type = "filter" self.type = "filter"
@ -73,28 +73,23 @@ class Pipeline:
if isinstance(body, str): if isinstance(body, str):
body = json.loads(body) body = json.loads(body)
if user and user.get("role", "admin") in self.valves.target_user_roles: model = body.get("model", "")
print(f"MODEL NAME: {model}")
# Get the content of the most recent message
user_message = get_last_user_message(body["messages"])
print("CURRENT MESSAGE:", user_message)
if model in self.valves.model_to_override:
messages = body.get("messages", []) messages = body.get("messages", [])
for message in messages: for message in messages:
if "images" in message: if "images" in message:
images.extend(message["images"]) images.extend(message["images"])
# Get the content of the most recent message
if messages:
user_message = get_last_user_message(body["messages"])
print("CURRENT MESSAGE:", user_message)
# Process the images with LLava
if images:
print("IMAGES: True") print("IMAGES: True")
llava_response = await self.process_images_with_llava(images, user_message, self.valves.vision_model,self.valves.ollama_base_url) llava_response = await self.process_images_with_llava(images, user_message, self.valves.vision_model,self.valves.ollama_base_url)
print("LLAVA RESPONSE:", llava_response)
# Override the content for the user role
for message in messages:
if message.get("role") == "user":
message["content"] = llava_response message["content"] = llava_response
print("LLAVA RESPONSE:", llava_response)
message.pop("images", None) # This will safely remove the 'images' key if it exists
else: else:
print("IMAGES: False") print("IMAGES: False")