vision-agent 0.2.66__py3-none-any.whl → 0.2.67__py3-none-any.whl

vision_agent/lmm/__init__.py

@@ -1 +1 @@
-from .lmm import LMM, AzureOpenAILMM, Message, OpenAILMM
+from .lmm import LMM, AzureOpenAILMM, Message, OllamaLMM, OpenAILMM

vision_agent/lmm/lmm.py

@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Union, cast
 
+import requests
 from openai import AzureOpenAI, OpenAI
 
 import vision_agent.tools as T
@@ -267,3 +268,84 @@ class AzureOpenAILMM(OpenAILMM):
         if json_mode:
             kwargs["response_format"] = {"type": "json_object"}
         self.kwargs = kwargs
+
+
+class OllamaLMM(LMM):
+    r"""An LMM class for the ollama."""
+
+    def __init__(
+        self,
+        model_name: str = "llava",
+        base_url: Optional[str] = "http://localhost:11434/api",
+        json_mode: bool = False,
+        **kwargs: Any,
+    ):
+        self.url = base_url
+        self.model_name = model_name
+        self.json_mode = json_mode
+        self.stream = False
+
+    def __call__(
+        self,
+        input: Union[str, List[Message]],
+    ) -> str:
+        if isinstance(input, str):
+            return self.generate(input)
+        return self.chat(input)
+
+    def chat(
+        self,
+        chat: List[Message],
+    ) -> str:
+        """Chat with the LMM model.
+
+        Parameters:
+            chat (List[Dict[str, str]]): A list of dictionaries containing the chat
+                messages. The messages can be in the format:
+                [{"role": "user", "content": "Hello!"}, ...]
+                or if it contains media, it should be in the format:
+                [{"role": "user", "content": "Hello!", "media": ["image1.jpg", ...]}, ...]
+        """
+        fixed_chat = []
+        for message in chat:
+            if "media" in message:
+                message["images"] = [encode_image(m) for m in message["media"]]
+                del message["media"]
+            fixed_chat.append(message)
+        url = f"{self.url}/chat"
+        model = self.model_name
+        messages = fixed_chat
+        data = {"model": model, "messages": messages, "stream": self.stream}
+        json_data = json.dumps(data)
+        response = requests.post(url, data=json_data)
+        if response.status_code != 200:
+            raise ValueError(f"Request failed with status code {response.status_code}")
+        response = response.json()
+        return response["message"]["content"]  # type: ignore
+
+    def generate(
+        self,
+        prompt: str,
+        media: Optional[List[Union[str, Path]]] = None,
+    ) -> str:
+
+        url = f"{self.url}/generate"
+        data = {
+            "model": self.model_name,
+            "prompt": prompt,
+            "images": [],
+            "stream": self.stream,
+        }
+
+        json_data = json.dumps(data)
+        if media and len(media) > 0:
+            for m in media:
+                data["images"].append(encode_image(m))  # type: ignore
+
+        response = requests.post(url, data=json_data)
+
+        if response.status_code != 200:
+            raise ValueError(f"Request failed with status code {response.status_code}")
+
+        response = response.json()
+        return response["response"]  # type: ignore

{vision_agent-0.2.66.dist-info → vision_agent-0.2.67.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.66
+Version: 0.2.67
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai

{vision_agent-0.2.66.dist-info → vision_agent-0.2.67.dist-info}/RECORD

@@ -5,8 +5,8 @@ vision_agent/agent/vision_agent.py,sha256=HC63BP4jPiR4lJLEkKQ-zMV5C5JwjnuZvc7hVj
 vision_agent/agent/vision_agent_prompts.py,sha256=jpGJjrxDrxZej5SSgsTEK1sSYttgkTiZqxZAU1jWfvk,8656
 vision_agent/fonts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vision_agent/fonts/default_font_ch_en.ttf,sha256=1YM0Z3XqLDjSNbF7ihQFSAIUdjF9m1rtHiNC_6QosTE,1594400
-vision_agent/lmm/__init__.py,sha256=3ro5lCIoS3DgEghOy0SPFrEhYvFnWZpVC5S5kSnIx6A,57
-vision_agent/lmm/lmm.py,sha256=ihmLYL_291HnELyMtfFKTCnPWnmuoEH2DDFmc4ynMG8,8945
+vision_agent/lmm/__init__.py,sha256=bw24xyQJHGzmph5e-bKCiTh9AX6tRFI2OUd0mofxjZI,68
+vision_agent/lmm/lmm.py,sha256=V7jfU94HwA-SiQLY14USHrSGtagVKCNGjZhW5MyKipo,11547
 vision_agent/tools/__init__.py,sha256=aE1O8cMeLDPO50Sc-CuAQ_Akh0viz7vBxDcVeZNqsA0,1604
 vision_agent/tools/prompts.py,sha256=V1z4YJLXZuUl_iZ5rY0M5hHc_2tmMEUKr0WocXKGt4E,1430
 vision_agent/tools/tool_utils.py,sha256=wzRacbUpqk9hhfX_Y08rL8qP0XCN2w-8IZoYLi3Upn4,869
@@ -17,7 +17,7 @@ vision_agent/utils/image_utils.py,sha256=_cdiS5YrLzqkq_ZgFUO897m5M4_SCIThwUy4lOk
 vision_agent/utils/sim.py,sha256=ci6Eta73dDgLP1Ajtknbgmf1g8aAvBHqlVQvBuLMKXQ,4427
 vision_agent/utils/type_defs.py,sha256=BlI8ywWHAplC7kYWLvt4AOdnKpEW3qWEFm-GEOSkrFQ,1792
 vision_agent/utils/video.py,sha256=rNmU9KEIkZB5-EztZNlUiKYN0mm_55A_2VGUM0QpqLA,8779
-vision_agent-0.2.66.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vision_agent-0.2.66.dist-info/METADATA,sha256=ou2H7-t3sn8KBwg6x56HCF1lf2epJi3-8-Doht3_Ci0,8363
-vision_agent-0.2.66.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
-vision_agent-0.2.66.dist-info/RECORD,,
+vision_agent-0.2.67.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vision_agent-0.2.67.dist-info/METADATA,sha256=BZKENJv_iaNU-XDqc5z4Ygx7k2jR4_7BbIdGoJE3voA,8363
+vision_agent-0.2.67.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+vision_agent-0.2.67.dist-info/RECORD,,