folder-classifier 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

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/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/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.dist-info → folder_classifier-0.3.4.dist-info}/METADATA

@@ -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.4.dist-info/RECORD

@@ -0,0 +1,10 @@
+folder_classifier/__init__.py,sha256=k0YWZyUNe7myJiKeX0OaXtJ30_3EGE-vsZiAUbqa-3E,46
+folder_classifier/app.py,sha256=_BUFSR7TJMoylrEwImDlln98C4BciPXoPqM6SQYkc1g,1255
+folder_classifier/classifier.py,sha256=XTqCfUAc4-vX1bS4LixPFNMgvNNFX2mLdWXekHHB-fQ,8285
+folder_classifier/deploy.py,sha256=06UAxz40IaP28e_RRohJoFwPUrWTaMquGbDylI-oHWA,424
+folder_classifier/dto.py,sha256=KRHEfXKlsIYkBySwHOlo4k5Z4J_F3eR3PeEdUAwj8lI,819
+folder_classifier/util.py,sha256=-Ptxkba5UzmhLrqoiiKZS3G56_cuAMkWlTyHcqdJkg0,3160
+folder_classifier-0.3.4.dist-info/METADATA,sha256=YwlhWYE_vMFYfBb0lib6kHUGZjUyH3soNKQiJq0jVmY,392
+folder_classifier-0.3.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+folder_classifier-0.3.4.dist-info/top_level.txt,sha256=36ugc9pEbNQ-mnzz4Ot2WVjY3t_LzAN6XOCjDFP4p4k,18
+folder_classifier-0.3.4.dist-info/RECORD,,

folder_classifier-0.3.3.dist-info/RECORD

@@ -1,10 +0,0 @@
-folder_classifier/__init__.py,sha256=k0YWZyUNe7myJiKeX0OaXtJ30_3EGE-vsZiAUbqa-3E,46
-folder_classifier/app.py,sha256=Xp3qkxQDWhFJKINMSMiDPtn1EEJrGm7XdTJAPZz6I4w,1358
-folder_classifier/classifier.py,sha256=mBt6QvWIWiKdfwqRXog_p0nRqEnZ6mKwgGf8VADx930,8146
-folder_classifier/deploy.py,sha256=06UAxz40IaP28e_RRohJoFwPUrWTaMquGbDylI-oHWA,424
-folder_classifier/dto.py,sha256=bVai09FuPXktOo5Y_9dzXXj0ZvriDwZOho0ZYmWvc7E,803
-folder_classifier/util.py,sha256=-Ptxkba5UzmhLrqoiiKZS3G56_cuAMkWlTyHcqdJkg0,3160
-folder_classifier-0.3.3.dist-info/METADATA,sha256=Tj220ru3uyno5el6TPyEJ12ew8pDzzdA-hIHKkD7Y8Y,392
-folder_classifier-0.3.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-folder_classifier-0.3.3.dist-info/top_level.txt,sha256=36ugc9pEbNQ-mnzz4Ot2WVjY3t_LzAN6XOCjDFP4p4k,18
-folder_classifier-0.3.3.dist-info/RECORD,,