Merge branch 'main' of https://github.com/santiagxf/llama-stack into santiagxf/azure-ai-inference

llama_stack/providers/utils/datasetio/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/providers/utils/datasetio/url_utils.py

@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import base64
+import io
+from urllib.parse import unquote
+
+import pandas
+
+from llama_models.llama3.api.datatypes import URL
+
+from llama_stack.providers.utils.memory.vector_store import parse_data_url
+
+
+def get_dataframe_from_url(url: URL):
+    df = None
+    if url.uri.endswith(".csv"):
+        df = pandas.read_csv(url.uri)
+    elif url.uri.endswith(".xlsx"):
+        df = pandas.read_excel(url.uri)
+    elif url.uri.startswith("data:"):
+        parts = parse_data_url(url.uri)
+        data = parts["data"]
+        if parts["is_base64"]:
+            data = base64.b64decode(data)
+        else:
+            data = unquote(data)
+            encoding = parts["encoding"] or "utf-8"
+            data = data.encode(encoding)
+
+        mime_type = parts["mimetype"]
+        mime_category = mime_type.split("/")[0]
+        data_bytes = io.BytesIO(data)
+
+        if mime_category == "text":
+            df = pandas.read_csv(data_bytes)
+        else:
+            df = pandas.read_excel(data_bytes)
+    else:
+        raise ValueError(f"Unsupported file type: {url}")
+
+    return df

llama_stack/providers/utils/inference/model_registry.py

@@ -4,11 +4,11 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Dict, List
+from typing import Dict
 
 from llama_models.sku_list import resolve_model
 
-from llama_stack.providers.datatypes import ModelDef, ModelsProtocolPrivate
+from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 
 
 class ModelRegistryHelper(ModelsProtocolPrivate):
@@ -28,14 +28,8 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
 
         return self.stack_to_provider_models_map[identifier]
 
-    async def register_model(self, model: ModelDef) -> None:
+    async def register_model(self, model: Model) -> None:
         if model.identifier not in self.stack_to_provider_models_map:
             raise ValueError(
                 f"Unsupported model {model.identifier}. Supported models: {self.stack_to_provider_models_map.keys()}"
             )
-
-    async def list_models(self) -> List[ModelDef]:
-        models = []
-        for llama_model, provider_model in self.stack_to_provider_models_map.items():
-            models.append(ModelDef(identifier=llama_model, llama_model=llama_model))
-        return models