Extending Connectors

1
class CustomSourceConnector(SourceConnector):
2
    async def validate_config(self):
3
        """
4
        Validate the connector configuration
5
        """
6
        if not self.config.api_key:
7
            raise ValueError("API key is required")
8
9
    async def test_connection(self):
10
        """
11
        Test the connection to your data source
12
        """
13
        try:
14
            await self.fetch_data(limit=1)
15
            return True
16
        except Exception as e:
17
            raise ConnectionError(f"Failed to connect: {str(e)}")
18
19
    async def get_schema(self):
20
        """
21
        Return the schema of your data source
22
        """
23
        return {
24
            "type": "object",
25
            "properties": {
26
                "content": {"type": "string"},
27
                "metadata": {
28
                    "type": "object",
29
                    "properties": {
30
                        "id": {"type": "string"},
31
                        "type": {"type": "string"}
32
                    }
33
                }
34
            }
35
        }

1
from airweave.connectors import DestinationConnector
2
from airweave.types import Entity, DestinationConfig
3
4
class CustomDestinationConnector(DestinationConnector):
5
    def __init__(self, config: DestinationConfig):
6
        super().__init__(config)
7
        self.config = config
8
9
    async def store_entities(self, entities: list[Entity]):
10
        """
11
        Store entities in your destination system
12
        """
13
        # Your custom storage logic here
14
        for entity in entities:
15
            await self.store_entity(entity)
16
17
    async def store_entity(self, entity: Entity):
18
        """
19
        Store a single entity
20
        """
21
        # Example implementation
22
        pass

1
class CustomDestinationConnector(DestinationConnector):
2
    async def initialize(self):
3
        """
4
        Initialize your destination system
5
        """
6
        # Example: Create necessary indexes
7
        pass
8
9
    async def cleanup(self):
10
        """
11
        Clean up resources
12
        """
13
        # Example: Close connections
14
        pass
15
16
    async def query(self, query: str, filter: dict = None):
17
        """
18
        Implement custom query logic
19
        """
20
        # Your custom query implementation
21
        pass

1
# sources/api/rest.py
2
from airweave.connectors import SourceConnector
3
from airweave.types import Entity
4
5
class RestApiSource(SourceConnector):
6
    async def generate_entities(self):
7
        async with aiohttp.ClientSession() as session:
8
            async with session.get(self.config.url) as response:
9
                data = await response.json()
10
11
                for item in data:
12
                    yield Entity(
13
                        content=item["content"],
14
                        metadata={
15
                            "source": "rest-api",
16
                            "id": item["id"]
17
                        }
18
                    )

1
# destinations/vector_stores/custom_store.py
2
from airweave.connectors import DestinationConnector
3
4
class CustomVectorStore(DestinationConnector):
5
    async def store_entities(self, entities):
6
        vectors = await self.get_embeddings(entities)
7
8
        async with self.client as client:
9
            for entity, vector in zip(entities, vectors):
10
                await client.upsert(
11
                    vector=vector,
12
                    metadata=entity.metadata
13
                )