gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

quick fix for crewai

Browse source
This commit is contained in:
hajdul88 2025-04-23 21:27:53 +02:00
parent 8a5bd3a826
commit 484520eba5
3 changed files with 101 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
examples/python/latest_ai_development/src/latest_ai_development/crew.py
View file

@ -1,13 +1,14 @@
from crewai import Agent, Crew, Process, Task from crewai import Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task from crewai.project import CrewBase, agent, crew, task, before_kickoff
from .tools import CogneeAdd, CogneeSearch from .tools import CogneeAdd, CogneeSearch
from crewai_tools import ( from crewai_tools import DirectoryReadTool
DirectoryReadTool
)
import os import os
docs_tool = DirectoryReadTool(directory='/Users/vasilije/cognee/examples/python/latest_ai_development/src/latest_ai_development/multimedia') # Determine multimedia input directory (can be overridden via env var)
multimedia_dir = os.getenv("MULTIMEDIA_DIR", os.path.join(os.path.dirname(__file__), "multimedia"))
docs_tool = DirectoryReadTool(directory=multimedia_dir)
# Utility function to format paths with file:// prefix # Utility function to format paths with file:// prefix
def format_file_paths(paths): def format_file_paths(paths):
@ -22,10 +23,14 @@ def format_file_paths(paths):
""" """
if isinstance(paths, str): if isinstance(paths, str):
# Split the paths by newline if it's a string output # Split the paths by newline if it's a string output
file_list = [line for line in paths.split('\n') if line.strip()] file_list = [line for line in paths.split("\n") if line.strip()]
# Format each path with file:// prefix # Format each path with file:// prefix
formatted_paths = [f"file://{os.path.abspath(path.strip())}" for path in file_list if "File paths:" not in path] formatted_paths = [
return '\n'.join(formatted_paths) f"file://{os.path.abspath(path.strip())}"
for path in file_list
if "File paths:" not in path
]
return "\n".join(formatted_paths)
return paths return paths
@ -55,13 +60,17 @@ class LatestAiDevelopment:
documents_cognee_add = CogneeAdd() documents_cognee_add = CogneeAdd()
documents_cognee_add.default_nodeset = ["documents"] documents_cognee_add.default_nodeset = ["documents"]
documents_cognee_add.name = "Add Documents to Memory" documents_cognee_add.name = "Add Documents to Memory"
documents_cognee_add.description = "Add document content to Cognee memory with documents nodeset" documents_cognee_add.description = (
"Add document content to Cognee memory with documents nodeset"
)
# CogneeAdd for reasoning/analysis with a "reasoning" nodeset # CogneeAdd for reasoning/analysis with a "reasoning" nodeset
reasoning_cognee_add = CogneeAdd() reasoning_cognee_add = CogneeAdd()
reasoning_cognee_add.default_nodeset = ["reasoning"] reasoning_cognee_add.default_nodeset = ["reasoning"]
reasoning_cognee_add.name = "Add Reasoning to Memory" reasoning_cognee_add.name = "Add Reasoning to Memory"
reasoning_cognee_add.description = "Add reasoning and analysis text to Cognee memory with reasoning nodeset" reasoning_cognee_add.description = (
"Add reasoning and analysis text to Cognee memory with reasoning nodeset"
)
# Create a wrapper for the DirectoryReadTool that formats output # Create a wrapper for the DirectoryReadTool that formats output
class FormattedDirectoryReadTool(DirectoryReadTool): class FormattedDirectoryReadTool(DirectoryReadTool):
@ -69,12 +78,13 @@ class LatestAiDevelopment:
result = super().__call__(*args, **kwargs) result = super().__call__(*args, **kwargs)
return format_file_paths(result) return format_file_paths(result)
formatted_docs_tool = FormattedDirectoryReadTool(directory='/Users/vasilije/cognee/examples/python/latest_ai_development/src/latest_ai_development/multimedia') # Use the project-local multimedia directory
formatted_docs_tool = FormattedDirectoryReadTool(directory=multimedia_dir)
return Agent( return Agent(
config=self.agents_config["researcher"], config=self.agents_config["researcher"],
tools=[formatted_docs_tool, documents_cognee_add, reasoning_cognee_add, cognee_search], tools=[formatted_docs_tool, documents_cognee_add, reasoning_cognee_add, cognee_search],
verbose=True verbose=True,
) )
@agent @agent
@ -107,12 +117,19 @@ class LatestAiDevelopment:
def reporting_task(self) -> Task: def reporting_task(self) -> Task:
return Task(config=self.tasks_config["reporting_task"], output_file="report.md") return Task(config=self.tasks_config["reporting_task"], output_file="report.md")
@before_kickoff
def dump_env(self, *args, **kwargs):
"""Print environment variables at startup."""
print("=== Environment Variables ===")
for key in sorted(os.environ):
print(f"{key}={os.environ[key]}")
@crew @crew
def crew(self) -> Crew: def crew(self) -> Crew:
"""Creates the LatestAiDevelopment crew""" """Creates the LatestAiDevelopment crew"""
# To learn how to add knowledge sources to your crew, check out the documentation: # To learn how to add knowledge sources to your crew, check out the documentation:
# https://docs.crewai.com/concepts/knowledge#what-is-knowledge # https://docs.crewai.com/concepts/knowledge#what-is-knowledge
print(self.tasks)
return Crew( return Crew(
agents=self.agents, # Automatically created by the @agent decorator agents=self.agents, # Automatically created by the @agent decorator
tasks=self.tasks, # Automatically created by the @task decorator tasks=self.tasks, # Automatically created by the @task decorator

1
examples/python/latest_ai_development/src/latest_ai_development/multimedia/dummy.txt Normal file
View file

@ -0,0 +1 @@
This is a dummy text file for testing DirectoryReadTool.

67
examples/python/latest_ai_development/src/latest_ai_development/tools/custom_tool.py
View file

@ -1,9 +1,7 @@
from crewai.tools import BaseTool from crewai.tools import BaseTool
from typing import Type, Dict, Any, List, Optional, Union from typing import Type, List, Optional
from pydantic import BaseModel, Field, root_validator from pydantic import BaseModel, Field, root_validator
from cognee.api.v1.search import SearchType from cognee.api.v1.search import SearchType
from cognee.modules.engine.models.Entity import Entity
from cognee.api.v1.search import SearchType
from cognee.modules.users.methods import get_default_user from cognee.modules.users.methods import get_default_user
from cognee.modules.pipelines import run_tasks, Task from cognee.modules.pipelines import run_tasks, Task
from cognee.tasks.experimental_tasks.node_set_edge_association import node_set_edge_association from cognee.tasks.experimental_tasks.node_set_edge_association import node_set_edge_association
@ -13,9 +11,18 @@ class CogneeAddInput(BaseModel):
"""Input schema for CogneeAdd tool.""" """Input schema for CogneeAdd tool."""
context: Optional[str] = Field(None, description="The text content to add to Cognee memory.") context: Optional[str] = Field(None, description="The text content to add to Cognee memory.")
file_paths: Optional[List[str]] = Field(None, description="List of file paths to add to Cognee memory.") file_paths: Optional[List[str]] = Field(
text: Optional[str] = Field(None, description="Alternative field for text content (maps to context).") None, description="List of file paths to add to Cognee memory."
reasoning: Optional[str] = Field(None, description="Alternative field for reasoning text (maps to context).") )
files: Optional[List[str]] = Field(
None, description="Alias for file_paths; list of file URLs or paths to add to memory."
)
text: Optional[str] = Field(
None, description="Alternative field for text content (maps to context)."
)
reasoning: Optional[str] = Field(
None, description="Alternative field for reasoning text (maps to context)."
)
node_set: List[str] = Field( node_set: List[str] = Field(
default=["default"], description="The list of node sets to store the data in." default=["default"], description="The list of node sets to store the data in."
) )
@ -23,16 +30,24 @@ class CogneeAddInput(BaseModel):
@root_validator(pre=True) @root_validator(pre=True)
def normalize_inputs(cls, values): def normalize_inputs(cls, values):
"""Normalize different input formats to standard fields.""" """Normalize different input formats to standard fields."""
# Map alias 'files' to 'file_paths' if provided
if values.get("files") and not values.get("file_paths"):
values["file_paths"] = values.get("files")
# Map text or reasoning to context if provided # Map text or reasoning to context if provided
if values.get('text') and not values.get('context'): if values.get("text") and not values.get("context"):
values['context'] = values.get('text') values["context"] = values.get("text")
if values.get('reasoning') and not values.get('context'): if values.get("reasoning") and not values.get("context"):
values['context'] = values.get('reasoning') values["context"] = values.get("reasoning")
# Map report_section to context if provided
if values.get("report_section") and not values.get("context"):
values["context"] = values.get("report_section")
# Validate that at least one input field is provided # Validate that at least one input field is provided
if not values.get('context') and not values.get('file_paths'): if not values.get("context") and not values.get("file_paths"):
raise ValueError("Either 'context', 'text', 'reasoning', or 'file_paths' must be provided") raise ValueError(
"Either 'context', 'text', 'reasoning', or 'file_paths' must be provided"
)
return values return values
@ -78,7 +93,9 @@ class CogneeAdd(BaseTool):
pipeline = run_tasks(tasks=tasks, user=user) pipeline = run_tasks(tasks=tasks, user=user)
async for pipeline_status in pipeline: async for pipeline_status in pipeline:
print(f"Pipeline run status: {pipeline_status.pipeline_name} - {pipeline_status.status}") print(
f"Pipeline run status: {pipeline_status.pipeline_name} - {pipeline_status.status}"
)
return run return run
except Exception as e: except Exception as e:
@ -91,7 +108,6 @@ class CogneeAdd(BaseTool):
# If loop is already running, create a new one # If loop is already running, create a new one
loop = asyncio.new_event_loop() loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop) asyncio.set_event_loop(loop)
result = loop.run_until_complete(main(node_set)) result = loop.run_until_complete(main(node_set))
return result.__name__ if hasattr(result, "__name__") else str(result) return result.__name__ if hasattr(result, "__name__") else str(result)
except Exception as e: except Exception as e:
@ -119,23 +135,23 @@ class CogneeSearchInput(BaseModel):
"""Normalize different input formats to standard fields.""" """Normalize different input formats to standard fields."""
# If the dictionary is empty, use a default query # If the dictionary is empty, use a default query
if not values: if not values:
values['query_text'] = "Latest AI developments" values["query_text"] = "Latest AI developments"
return values return values
# Map alternative search fields to query_text # Map alternative search fields to query_text
if values.get('query') and not values.get('query_text'): if values.get("query") and not values.get("query_text"):
values['query_text'] = values.get('query') values["query_text"] = values.get("query")
if values.get('search_term') and not values.get('query_text'): if values.get("search_term") and not values.get("query_text"):
values['query_text'] = values.get('search_term') values["query_text"] = values.get("search_term")
# If security_context is provided but no query, use a default # If security_context is provided but no query, use a default
if 'security_context' in values and not values.get('query_text'): if "security_context" in values and not values.get("query_text"):
values['query_text'] = "Latest AI developments" values["query_text"] = "Latest AI developments"
# Ensure query_text is present # Ensure query_text is present
if not values.get('query_text'): if not values.get("query_text"):
values['query_text'] = "Latest AI developments" values["query_text"] = "Latest AI developments"
return values return values
@ -168,10 +184,11 @@ class CogneeSearch(BaseTool):
async def main(query, ns): async def main(query, ns):
try: try:
# Use 'datasets' to specify which node sets (datasets) to search
result = await cognee.search( result = await cognee.search(
query_type=SearchType.GRAPH_COMPLETION,
query_text=query + " Only return results from context", query_text=query + " Only return results from context",
node_set=ns # Pass the node_set to the search query_type=SearchType.GRAPH_COMPLETION,
datasets=ns,
) )
return result return result
except Exception as e: except Exception as e: