Move off v1

This commit is contained in:
Brandon Hancock
2024-09-03 15:57:29 -04:00
parent d19bba72b0
commit 35fe222ca1
39 changed files with 752 additions and 550 deletions

View File

@@ -1,26 +1,32 @@
from typing import Any, Optional, Type
from embedchain.models.data_type import DataType
from pydantic.v1 import BaseModel, Field
from pydantic import BaseModel, Field
from ..rag.rag_tool import RagTool
class FixedDOCXSearchToolSchema(BaseModel):
"""Input for DOCXSearchTool."""
docx: Optional[str] = Field(..., description="Mandatory docx path you want to search")
docx: Optional[str] = Field(
..., description="Mandatory docx path you want to search"
)
search_query: str = Field(
...,
description="Mandatory search query you want to use to search the DOCX's content",
)
class DOCXSearchToolSchema(FixedDOCXSearchToolSchema):
"""Input for DOCXSearchTool."""
search_query: str = Field(
...,
description="Mandatory search query you want to use to search the DOCX's content",
)
class DOCXSearchTool(RagTool):
name: str = "Search a DOCX's content"
description: str = (
@@ -56,9 +62,9 @@ class DOCXSearchTool(RagTool):
self,
**kwargs: Any,
) -> Any:
search_query = kwargs.get('search_query')
search_query = kwargs.get("search_query")
if search_query is None:
search_query = kwargs.get('query')
search_query = kwargs.get("query")
docx = kwargs.get("docx")
if docx is not None: