folder-classifier 0.3.3__tar.gz → 0.3.4__tar.gz

{folder_classifier-0.3.3 → folder_classifier-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: folder-classifier
-Version: 0.3.3
+Version: 0.3.4
 Summary: Deploy folder classifier API to a Ray cluster
 Author: Crispin Almodovar
 Author-email: 

{folder_classifier-0.3.3 → folder_classifier-0.3.4}/folder_classifier/app.py

@@ -16,10 +16,8 @@ class FolderClassifierAPI:
         assert model_config.app_name and model_config.deployment, "Invalid ModelConfig values"
         logging.basicConfig(level=logging.INFO)
         self.logger = logging.getLogger(__name__)
-        self.logger.info(f"Initializing Folder Classifier model: {model_config}")
-        model_handle = serve.get_deployment_handle(app_name=model_config.app_name, deployment_name=model_config.deployment)
-        self.classifier = FolderClassifier(model_handle)
-        self.logger.info(f"Successfully initialized Folder Classifier API")
+        self.classifier = FolderClassifier(app_name=model_config.app_name, deployment=model_config.deployment, model=model_config.model)
+        self.logger.info(f"Successfully initialized Folder Classifier using config: {model_config}")
 
     @web_api.post("/predict")
     async def predict(self, request: FolderClassificationRequest) -> FolderClassificationResponse:

{folder_classifier-0.3.3 → folder_classifier-0.3.4}/folder_classifier/classifier.py

@@ -1,7 +1,9 @@
+import json
 import logging
 import os
 from typing import Tuple, Dict, Any
 
+from ray import serve
 from ray.serve.handle import DeploymentHandle
 
 from folder_classifier.dto import FolderClassificationRequest, FolderClassification
@@ -81,18 +83,19 @@ FOLDER_CLASSIFICATION_SCHEMA = FolderClassification.model_json_schema()
 
 
 class FolderClassifier:
-    def __init__(self, model_handle: DeploymentHandle):
+    def __init__(self, app_name: str, deployment: str, model: str):
         self.logger = logging.getLogger(__name__)
-        self.model_handle = model_handle
-        self.llm = AsyncOpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)
+        self.model_handle = serve.get_deployment_handle(app_name=app_name, deployment_name=deployment)
+        self.model = model
         self.logger.info(f"Successfully initialized FolderClassifier")
 
     async def predict(self, request: FolderClassificationRequest) -> Tuple[str, str]:
         content = ""
         try:
             chat_completion_request = self._to_chat_completion_request(request)
-            response = await self.llm.chat.completions.create(**chat_completion_request)
-            content = response.choices[0].message.content
+            response = await self.model_handle.create_chat_completion_internal.remote(chat_completion_request)
+            response_dict = json.loads(response.body)
+            content = response_dict["choices"][0]["message"]["content"]
             result = FolderClassification.model_validate_json(content)
         except Exception as ex:
             self.logger.error(f"Failed to parse response: {content}\n{ex}")
@@ -102,13 +105,12 @@ class FolderClassifier:
                 result = FolderClassification(category="other", reasoning="NA")
         return result.category, result.reasoning
 
-    @staticmethod
-    def _to_chat_completion_request(request: FolderClassificationRequest) -> Dict[str, Any]:
+    def _to_chat_completion_request(self, request: FolderClassificationRequest) -> Dict[str, Any]:
         input_paths = request.items
         folder = build_folder(input_paths)
         folder_tree = render_tree(folder)
         chat_completion_request = {
-            "model": MODEL,
+            "model": self.model,
             "messages": [
                 {"role": "system", "content": SYSTEM_PROMPT},
                 {"role": "user", "content": USER_PROMPT_TEMPLATE.replace("{folder_tree}", folder_tree)}

{folder_classifier-0.3.3 → folder_classifier-0.3.4}/folder_classifier/dto.py

@@ -5,6 +5,7 @@ from pydantic import BaseModel, Field, ConfigDict
 class ModelConfig(BaseModel):
     app_name: str
     deployment: str
+    model: str
 
 
 class AppConfig(BaseModel):

{folder_classifier-0.3.3 → folder_classifier-0.3.4}/folder_classifier.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: folder-classifier
-Version: 0.3.3
+Version: 0.3.4
 Summary: Deploy folder classifier API to a Ray cluster
 Author: Crispin Almodovar
 Author-email: 

{folder_classifier-0.3.3 → folder_classifier-0.3.4}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = folder-classifier
-version = 0.3.3
+version = 0.3.4
 author = Crispin Almodovar
 author_email = 
 description = Deploy folder classifier API to a Ray cluster