put a flag

This commit is contained in:
Lorenze Jay
2024-11-19 17:42:44 -08:00
parent 4663997b4c
commit 76da972ce9
3 changed files with 14 additions and 9 deletions

View File

@@ -231,8 +231,8 @@ class Agent(BaseAgent):
self._validate_docker_installation() self._validate_docker_installation()
# Initialize the Knowledge object if knowledge_sources are provided # Initialize the Knowledge object if knowledge_sources are provided
if self.crew and self.crew.knowledge: if self.crew and self.crew.knowledge_store:
self._knowledge = self.crew.knowledge self._knowledge = self.crew.knowledge_store
else: else:
self._knowledge = None self._knowledge = None
@@ -282,8 +282,8 @@ class Agent(BaseAgent):
task_prompt += self.i18n.slice("memory").format(memory=memory) task_prompt += self.i18n.slice("memory").format(memory=memory)
# Integrate the knowledge base # Integrate the knowledge base
if self.crew and self.crew.knowledge: if self.crew and self.crew.knowledge_store and self._knowledge:
knowledge_snippets: List[Dict[str, Any]] = self.crew.knowledge.query( knowledge_snippets: List[Dict[str, Any]] = self.crew.knowledge_store.query(
[task.prompt()] [task.prompt()]
) )
if knowledge_snippets: if knowledge_snippets:

View File

@@ -199,9 +199,14 @@ class Crew(BaseModel):
default=None, default=None,
description="Knowledge sources for the agent.", description="Knowledge sources for the agent.",
) )
knowledge: Optional[Knowledge] = Field( knowledge_store: Optional[Knowledge] = Field(
default=None, description="Knowledge Source for the crew." default=None, description="Knowledge Source for the crew."
) )
knowledge: Optional[bool] = Field(
default=False,
description="Whether the crew should use knowledge to store memories of it's execution",
)
@field_validator("id", mode="before") @field_validator("id", mode="before")
@classmethod @classmethod
@@ -278,9 +283,10 @@ class Crew(BaseModel):
@model_validator(mode="after") @model_validator(mode="after")
def create_crew_knowledge(self) -> "Crew": def create_crew_knowledge(self) -> "Crew":
self.knowledge = Knowledge( if self.knowledge:
sources=self.knowledge_sources or [], embedder_config=self.embedder self.knowledge_store = Knowledge(
) sources=self.knowledge_sources or [], embedder_config=self.embedder
)
return self return self
@model_validator(mode="after") @model_validator(mode="after")

View File

@@ -416,7 +416,6 @@ def test_hybrid_string_and_files(mock_vector_db, tmpdir):
mock_vector_db.query.assert_called_once() mock_vector_db.query.assert_called_once()
@pytest.mark.vcr(filter_headers=["authorization"])
def test_pdf_knowledge_source(mock_vector_db): def test_pdf_knowledge_source(mock_vector_db):
# Get the directory of the current file # Get the directory of the current file
current_dir = Path(__file__).parent current_dir = Path(__file__).parent