Merge branch 'feat/support-extractor-tools' into deploy/dev

api/controllers/console/__init__.py

@@ -55,9 +55,9 @@
     datasets_document,
     datasets_segments,
     external,
+    fta_test,
     hit_testing,
     website,
-    fta_test,
 )
 
 # Import explore controllers

api/controllers/console/datasets/fta_test.py

@@ -1,30 +1,20 @@
 import json
 
-from flask import Response
 import requests
+from flask import Response
+from flask_restful import Resource, reqparse
 from sqlalchemy import text
 
 from controllers.console import api
-
-from flask_restful import Resource, reqparse
 from extensions.ext_database import db
-
-from models.fta import ComponentFailure, ComponentFailureStats
-
 from extensions.ext_storage import storage
+from models.fta import ComponentFailure, ComponentFailureStats
 
 
 class FATTestApi(Resource):
-
     def post(self):
         parser = reqparse.RequestParser()
-        parser.add_argument(
-            "log_process_data",
-            nullable=False,
-            required=True,
-            type=str,
-            location="args"
-        )
+        parser.add_argument("log_process_data", nullable=False, required=True, type=str, location="args")
         args = parser.parse_args()
         print(args["log_process_data"])
         # Extract the JSON string from the text field
@@ -47,7 +37,7 @@ def post(self):
                     FailureMode=data["FailureMode"],
                     Cause=data["Cause"],
                     RepairAction=data["RepairAction"],
-                    Technician=data["Technician"]
+                    Technician=data["Technician"],
                 )
                 db.session.add(component_failure)
                 db.session.commit()
@@ -98,66 +88,49 @@ def post(self):
         # Convert stats to list of tuples format
         stats_list = []
         for stat in component_failure_stats:
-            stats_list.append((
-                stat.StatID,
-                stat.Component,
-                stat.FailureMode,
-                stat.Cause,
-                stat.PossibleAction,
-                stat.Probability,
-                stat.MTBF
-            ))
+            stats_list.append(
+                (
+                    stat.StatID,
+                    stat.Component,
+                    stat.FailureMode,
+                    stat.Cause,
+                    stat.PossibleAction,
+                    stat.Probability,
+                    stat.MTBF,
+                )
+            )
         return {"data": stats_list}, 200
 
 
 # generate-fault-tree
 class GenerateFaultTreeApi(Resource):
-
     def post(self):
         parser = reqparse.RequestParser()
-        parser.add_argument(
-            "llm_text",
-            nullable=False,
-            required=True,
-            type=str,
-            location="args"
-        )
+        parser.add_argument("llm_text", nullable=False, required=True, type=str, location="args")
         args = parser.parse_args()
         entities = args["llm_text"].replace("```", "").replace("\\n", "\n")
         print(entities)
-        request_data = {
-            "fault_tree_text": entities
-        }
+        request_data = {"fault_tree_text": entities}
         url = "https://fta.cognitech-dev.live/generate-fault-tree"
-        headers = {
-            "accept": "application/json",
-            "Content-Type": "application/json"
-        }
+        headers = {"accept": "application/json", "Content-Type": "application/json"}
 
         response = requests.post(url, json=request_data, headers=headers)
         print(response.json())
         return {"data": response.json()}, 200
 
 
 class ExtractSVGApi(Resource):
-
     def post(self):
         parser = reqparse.RequestParser()
-        parser.add_argument(
-            "svg_text",
-            nullable=False,
-            required=True,
-            type=str,
-            location="args"
-        )
+        parser.add_argument("svg_text", nullable=False, required=True, type=str, location="args")
         args = parser.parse_args()
         # svg_text = ''.join(args["svg_text"].splitlines())
-        svg_text = args["svg_text"].replace('\n', '')
-        svg_text = svg_text.replace('\"', '"')
+        svg_text = args["svg_text"].replace("\n", "")
+        svg_text = svg_text.replace('"', '"')
         print(svg_text)
         svg_text_json = json.loads(svg_text)
         svg_content = svg_text_json.get("data").get("svg_content")[0]
-        svg_content = svg_content.replace('\n', '').replace('\"', '"')
+        svg_content = svg_content.replace("\n", "").replace('"', '"')
         file_key = "fta_svg/" + "fat.svg"
         if storage.exists(file_key):
             storage.delete(file_key)

api/core/file/file_manager.py

@@ -1,6 +1,6 @@
 import base64
-from pathlib import Path
 import tempfile
+from pathlib import Path
 
 from configs import dify_config
 from core.file import file_repository

api/core/tools/provider/builtin/file_extractor/file_extractor.py

@@ -1,7 +1,5 @@
 from typing import Any
 
-from core.tools.errors import ToolProviderCredentialValidationError
-from core.tools.provider.builtin.dalle.tools.dalle2 import DallE2Tool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
 
 

api/core/tools/provider/builtin/file_extractor/tools/file_extractor.py

@@ -1,12 +1,7 @@
-from base64 import b64decode
-import json
 import tempfile
 from typing import Any, Union
 
-from openai import OpenAI
-from yarl import URL
 from core.file.enums import FileType
-
 from core.file.file_manager import download_to_target_path
 from core.rag.extractor.text_extractor import TextExtractor
 from core.rag.splitter.fixed_text_splitter import FixedRecursiveCharacterTextSplitter
@@ -45,6 +40,6 @@ def _invoke(
 
                 content = "\n".join([chunk.page_content for chunk in chunks])
                 return self.create_text_message(content)
-                
+
         else:
             raise ToolParameterValidationError("Please provide either file")

api/models/fta.py

@@ -1,32 +1,10 @@
-import base64
-import enum
-import hashlib
-import hmac
-import json
-import logging
-import os
-import pickle
-import re
-import time
-from json import JSONDecodeError
-
-from sqlalchemy import func
-from sqlalchemy.dialects.postgresql import JSONB
-
-from configs import dify_config
-from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from extensions.ext_database import db
-from extensions.ext_storage import storage
-
-from .account import Account
-from .model import App, Tag, TagBinding, UploadFile
-from .types import StringUUID
 
 
 class ComponentFailure(db.Model):
     __tablename__ = "component_failure"
     __table_args__ = (
-        db.UniqueConstraint('Date', 'Component', 'FailureMode', 'Cause', 'Technician', name='unique_failure_entry'),
+        db.UniqueConstraint("Date", "Component", "FailureMode", "Cause", "Technician", name="unique_failure_entry"),
     )
 
     FailureID = db.Column(db.Integer, primary_key=True, autoincrement=True)