Update schema_utils imports

llama_stack/apis/agents/agents.py

@@ -19,7 +19,6 @@ from typing import (
     runtime_checkable,
 )
 
-from llama_models.schema_utils import json_schema_type, register_schema, webmethod
 from pydantic import BaseModel, ConfigDict, Field
 
 from llama_stack.apis.common.content_types import URL, ContentDelta, InterleavedContent
@@ -38,6 +37,7 @@ from llama_stack.apis.inference import (
 from llama_stack.apis.safety import SafetyViolation
 from llama_stack.apis.tools import ToolDef
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
 
 
 class Attachment(BaseModel):

llama_stack/apis/batch_inference/batch_inference.py

@@ -6,7 +6,6 @@
 
 from typing import List, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.inference import (
@@ -21,6 +20,7 @@ from llama_stack.apis.inference import (
     ToolDefinition,
     ToolPromptFormat,
 )
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 @json_schema_type

llama_stack/apis/benchmarks/benchmarks.py

@@ -5,10 +5,10 @@
 # the root directory of this source tree.
 from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.resource import Resource, ResourceType
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 class CommonBenchmarkFields(BaseModel):

llama_stack/apis/common/content_types.py

@@ -8,9 +8,10 @@ from enum import Enum
 from typing import Annotated, List, Literal, Optional, Union
 
 from llama_models.llama3.api.datatypes import ToolCall
-from llama_models.schema_utils import json_schema_type, register_schema
 from pydantic import BaseModel, Field, model_validator
 
+from llama_stack.schema_utils import json_schema_type, register_schema
+
 
 @json_schema_type
 class URL(BaseModel):

llama_stack/apis/common/deployment_types.py

@@ -7,10 +7,10 @@
 from enum import Enum
 from typing import Any, Dict, Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel
 
 from llama_stack.apis.common.content_types import URL
+from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type

llama_stack/apis/common/job_types.py

@@ -5,9 +5,10 @@
 # the root directory of this source tree.
 from enum import Enum
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class Job(BaseModel):

llama_stack/apis/common/training_types.py

@@ -7,9 +7,10 @@
 from datetime import datetime
 from typing import Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class PostTrainingMetric(BaseModel):

llama_stack/apis/common/type_system.py

@@ -6,10 +6,11 @@
 
 from typing import Literal, Union
 
-from llama_models.schema_utils import json_schema_type, register_schema
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
+from llama_stack.schema_utils import json_schema_type, register_schema
+
 
 @json_schema_type
 class StringType(BaseModel):

llama_stack/apis/datasetio/datasetio.py

@@ -6,10 +6,10 @@
 
 from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.datasets import Dataset
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 @json_schema_type

llama_stack/apis/datasets/datasets.py

@@ -6,12 +6,12 @@
 
 from typing import Any, Dict, List, Literal, Optional, Protocol
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.common.content_types import URL
 from llama_stack.apis.common.type_system import ParamType
 from llama_stack.apis.resource import Resource, ResourceType
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 class CommonDatasetFields(BaseModel):

llama_stack/apis/datatypes.py

@@ -6,7 +6,7 @@
 
 from enum import Enum
 
-from llama_models.schema_utils import json_schema_type
+from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type

llama_stack/apis/eval/eval.py

@@ -6,7 +6,6 @@
 
 from typing import Any, Dict, List, Literal, Optional, Protocol, Union
 
-from llama_models.schema_utils import json_schema_type, register_schema, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
@@ -15,6 +14,7 @@ from llama_stack.apis.common.job_types import Job, JobStatus
 from llama_stack.apis.inference import SamplingParams, SystemMessage
 from llama_stack.apis.scoring import ScoringResult
 from llama_stack.apis.scoring_functions import ScoringFnParams
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
 
 
 @json_schema_type

llama_stack/apis/inference/inference.py

@@ -25,7 +25,6 @@ from llama_models.llama3.api.datatypes import (
     ToolDefinition,
     ToolPromptFormat,
 )
-from llama_models.schema_utils import json_schema_type, register_schema, webmethod
 from pydantic import BaseModel, Field, field_validator
 from typing_extensions import Annotated
 
@@ -33,6 +32,7 @@ from llama_stack.apis.common.content_types import ContentDelta, InterleavedConte
 from llama_stack.apis.models import Model
 from llama_stack.apis.telemetry.telemetry import MetricResponseMixin
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
 
 
 class LogProbConfig(BaseModel):

llama_stack/apis/inspect/inspect.py

@@ -6,9 +6,10 @@
 
 from typing import List, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
+from llama_stack.schema_utils import json_schema_type, webmethod
+
 
 @json_schema_type
 class ProviderInfo(BaseModel):

llama_stack/apis/models/models.py

@@ -7,11 +7,11 @@
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, ConfigDict, Field
 
 from llama_stack.apis.resource import Resource, ResourceType
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 class CommonModelFields(BaseModel):

llama_stack/apis/post_training/post_training.py

@@ -8,13 +8,13 @@ from datetime import datetime
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Protocol, Union
 
-from llama_models.schema_utils import json_schema_type, register_schema, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
 from llama_stack.apis.common.content_types import URL
 from llama_stack.apis.common.job_types import JobStatus
 from llama_stack.apis.common.training_types import Checkpoint
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
 
 
 @json_schema_type

llama_stack/apis/safety/safety.py

@@ -7,12 +7,12 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.inference import Message
 from llama_stack.apis.shields import Shield
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 @json_schema_type

llama_stack/apis/scoring/scoring.py

@@ -6,10 +6,10 @@
 
 from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 # mapping of metric to value
 ScoringResultRow = Dict[str, Any]

llama_stack/apis/scoring_functions/scoring_functions.py

@@ -16,12 +16,12 @@ from typing import (
     runtime_checkable,
 )
 
-from llama_models.schema_utils import json_schema_type, register_schema, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
 from llama_stack.apis.common.type_system import ParamType
 from llama_stack.apis.resource import Resource, ResourceType
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
 
 
 # Perhaps more structure can be imposed on these functions. Maybe they could be associated

llama_stack/apis/shields/shields.py

@@ -6,11 +6,11 @@
 
 from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.resource import Resource, ResourceType
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 class CommonShieldFields(BaseModel):

llama_stack/apis/synthetic_data_generation/synthetic_data_generation.py

@@ -7,10 +7,10 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Protocol, Union
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.inference import Message
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 class FilteringFunction(Enum):

llama_stack/apis/telemetry/telemetry.py

@@ -18,10 +18,11 @@ from typing import (
 )
 
 from llama_models.llama3.api.datatypes import Primitive
-from llama_models.schema_utils import json_schema_type, register_schema, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
+
 # Add this constant near the top of the file, after the imports
 DEFAULT_TTL_DAYS = 7
 

llama_stack/apis/tools/rag_tool.py

@@ -7,12 +7,12 @@
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Union
 
-from llama_models.schema_utils import json_schema_type, register_schema, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated, Protocol, runtime_checkable
 
 from llama_stack.apis.common.content_types import URL, InterleavedContent
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
 
 
 @json_schema_type

llama_stack/apis/tools/tools.py

@@ -7,13 +7,13 @@
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 from typing_extensions import Protocol, runtime_checkable
 
 from llama_stack.apis.common.content_types import URL, InterleavedContent
 from llama_stack.apis.resource import Resource, ResourceType
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 from .rag_tool import RAGToolRuntime
 

llama_stack/apis/vector_dbs/vector_dbs.py

@@ -6,11 +6,11 @@
 
 from typing import List, Literal, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_stack.apis.resource import Resource, ResourceType
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 @json_schema_type

llama_stack/apis/vector_io/vector_io.py

@@ -10,12 +10,12 @@
 # the root directory of this source tree.
 from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
 
-from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.inference import InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
+from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 class Chunk(BaseModel):

llama_stack/providers/datatypes.py

@@ -7,7 +7,6 @@
 from typing import Any, List, Optional, Protocol
 from urllib.parse import urlparse
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
 from llama_stack.apis.benchmarks import Benchmark
@@ -18,6 +17,7 @@ from llama_stack.apis.scoring_functions import ScoringFn
 from llama_stack.apis.shields import Shield
 from llama_stack.apis.tools import Tool
 from llama_stack.apis.vector_dbs import VectorDB
+from llama_stack.schema_utils import json_schema_type
 
 
 class ModelsProtocolPrivate(Protocol):

llama_stack/providers/inline/inference/vllm/config.py

@@ -4,10 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field, field_validator
 
 from llama_stack.providers.utils.inference import supported_inference_models
+from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type

llama_stack/providers/inline/post_training/torchtune/post_training.py

@@ -6,8 +6,6 @@
 from datetime import datetime
 from typing import Any, Dict, Optional
 
-from llama_models.schema_utils import webmethod
-
 from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import Datasets
 from llama_stack.apis.post_training import (
@@ -27,6 +25,7 @@ from llama_stack.providers.inline.post_training.torchtune.config import (
 from llama_stack.providers.inline.post_training.torchtune.recipes.lora_finetuning_single_device import (
     LoraFinetuningSingleDevice,
 )
+from llama_stack.schema_utils import webmethod
 
 
 class TorchtunePostTrainingImpl:

llama_stack/providers/inline/vector_io/faiss/config.py

@@ -6,13 +6,13 @@
 
 from typing import Any, Dict
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel
 
 from llama_stack.providers.utils.kvstore.config import (
     KVStoreConfig,
     SqliteKVStoreConfig,
 )
+from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type

llama_stack/providers/remote/inference/cerebras/config.py

@@ -7,9 +7,10 @@
 import os
 from typing import Any, Dict, Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field, SecretStr
 
+from llama_stack.schema_utils import json_schema_type
+
 DEFAULT_BASE_URL = "https://api.cerebras.ai"
 
 

llama_stack/providers/remote/inference/databricks/config.py

@@ -5,9 +5,10 @@
 # the root directory of this source tree.
 
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class DatabricksImplConfig(BaseModel):

llama_stack/providers/remote/inference/fireworks/config.py

@@ -6,9 +6,10 @@
 
 from typing import Any, Dict, Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field, SecretStr
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class FireworksImplConfig(BaseModel):

llama_stack/providers/remote/inference/groq/config.py

@@ -6,9 +6,10 @@
 
 from typing import Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class GroqConfig(BaseModel):

llama_stack/providers/remote/inference/nvidia/config.py

@@ -7,9 +7,10 @@
 import os
 from typing import Any, Dict, Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field, SecretStr
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class NVIDIAConfig(BaseModel):

llama_stack/providers/remote/inference/runpod/config.py

@@ -6,9 +6,10 @@
 
 from typing import Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class RunpodImplConfig(BaseModel):

llama_stack/providers/remote/inference/sambanova/config.py

@@ -6,9 +6,10 @@
 
 from typing import Any, Dict, Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class SambaNovaImplConfig(BaseModel):

llama_stack/providers/remote/inference/tgi/config.py

@@ -6,9 +6,10 @@
 
 from typing import Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field, SecretStr
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class TGIImplConfig(BaseModel):

llama_stack/providers/remote/inference/together/config.py

@@ -6,9 +6,10 @@
 
 from typing import Any, Dict, Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field, SecretStr
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class TogetherImplConfig(BaseModel):

llama_stack/providers/remote/inference/vllm/config.py

@@ -6,9 +6,10 @@
 
 from typing import Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class VLLMInferenceAdapterConfig(BaseModel):

llama_stack/providers/remote/safety/bedrock/config.py

@@ -5,9 +5,8 @@
 # the root directory of this source tree.
 
 
-from llama_models.schema_utils import json_schema_type
-
 from llama_stack.providers.utils.bedrock.config import BedrockBaseConfig
+from llama_stack.schema_utils import json_schema_type
 
 
 @json_schema_type

llama_stack/providers/remote/vector_io/pgvector/config.py

@@ -4,9 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class PGVectorVectorIOConfig(BaseModel):

llama_stack/providers/remote/vector_io/qdrant/config.py

@@ -6,9 +6,10 @@
 
 from typing import Optional
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class QdrantVectorIOConfig(BaseModel):

llama_stack/providers/utils/kvstore/sqlite/config.py

@@ -4,9 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class SqliteControlPlaneConfig(BaseModel):