1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from manas_ai.nodes import QANode
from manas_ai.models import Document
from typing import Dict, Any, List

class CustomQANode(QANode):
    """Custom QA node with specialized processing."""
    
    async def process_question(
        self,
        question: str,
        context: str = None
    ) -> Dict[str, Any]:
        """Custom question processing logic."""
        # Add preprocessing
        enhanced_question = self.enhance_question(question)
        
        # Get base response
        response = await super().process_question(
            enhanced_question,
            context
        )
        
        # Add postprocessing
        response["enhanced"] = True
        response["original_question"] = question
        return response
    
    def enhance_question(self, question: str) -> str:
        """Add custom question enhancement."""
        # Add your enhancement logic
        return f"{question} (Enhanced)"

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from manas_ai.nodes import DocumentNode
from manas_ai.models import Document
from typing import List

class CustomDocumentNode(DocumentNode):
    """Custom document processing node."""
    
    async def process_document(self, doc: Document) -> List[Document]:
        """Custom document processing."""
        # Implement custom chunking
        chunks = self.custom_chunk_document(doc)
        
        # Process each chunk
        processed = []
        for chunk in chunks:
            # Add custom processing
            enhanced = self.enhance_chunk(chunk)
            processed.append(enhanced)
            
        return processed
    
    def custom_chunk_document(self, doc: Document) -> List[str]:
        """Custom document chunking strategy."""
        # Implement your chunking logic
        return doc.content.split("\n\n")
    
    def enhance_chunk(self, chunk: str) -> Document:
        """Enhance a document chunk."""
        # Add your enhancement logic
        return Document(
            content=f"Enhanced: {chunk}",
            metadata={"enhanced": True}
        )

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from manas_ai.vectorstores import BaseVectorStore
from manas_ai.models import Document
from typing import List, Tuple

class CustomVectorStore(BaseVectorStore):
    """Custom vector store implementation."""
    
    def __init__(self, dimension: int):
        super().__init__()
        self.dimension = dimension
        self.initialize_store()
    
    def initialize_store(self):
        """Initialize storage backend."""
        # Set up your storage
        self.vectors = []
        self.documents = []
    
    async def add_documents(
        self,
        documents: List[Document]
    ) -> int:
        """Add documents to store."""
        # Get embeddings
        embeddings = await self.get_embeddings(
            [doc.content for doc in documents]
        )
        
        # Store vectors and documents
        for vec, doc in zip(embeddings, documents):
            self.vectors.append(vec)
            self.documents.append(doc)
            
        return len(documents)
    
    async def similarity_search(
        self,
        query: str,
        k: int = 3
    ) -> List[Tuple[Document, float]]:
        """Search for similar documents."""
        # Get query embedding
        query_vec = await self.get_embeddings([query])[0]
        
        # Calculate similarities
        scores = self.calculate_similarities(query_vec)
        
        # Get top k results
        results = []
        for idx, score in self.get_top_k(scores, k):
            results.append((self.documents[idx], score))
            
        return results

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
from manas_ai.providers import BaseProvider
from manas_ai.models import ChatMessage
from typing import List, Dict, Any

class CustomProvider(BaseProvider):
    """Custom LLM provider implementation."""
    
    def __init__(self, **kwargs):
        super().__init__()
        self.api_key = kwargs.get("api_key")
        self.initialize_client()
    
    def initialize_client(self):
        """Initialize API client."""
        # Set up your API client
        self.client = YourAPIClient(self.api_key)
    
    async def generate(
        self,
        prompt: str,
        **kwargs
    ) -> str:
        """Generate text completion."""
        response = await self.client.complete(
            prompt,
            max_tokens=kwargs.get("max_tokens", 100),
            temperature=kwargs.get("temperature", 0.7)
        )
        return response.text
    
    async def chat(
        self,
        messages: List[ChatMessage],
        **kwargs
    ) -> str:
        """Generate chat completion."""
        formatted = self.format_messages(messages)
        response = await self.client.chat(
            formatted,
            **kwargs
        )
        return response.message

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from core import Tool
from typing import Dict, Any

def custom_tool_function(
    param1: str,
    param2: int = 10
) -> Dict[str, Any]:
    """
    Custom tool implementation.
    
    Args:
        param1: First parameter
        param2: Second parameter (default: 10)
        
    Returns:
        Tool execution results
    """
    result = your_implementation(param1, param2)
    return {
        "status": "success",
        "result": result
    }

# Create tool
custom_tool = Tool(
    name="custom_tool",
    description="A custom tool that does something",
    function=custom_tool_function,
    metadata={
        "parameters": {
            "param1": {
                "type": "string",
                "description": "First parameter"
            },
            "param2": {
                "type": "integer",
                "description": "Second parameter",
                "default": 10
            }
        },
        "returns": {
            "type": "object",
            "description": "Tool results"
        }
    }
)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import pytest
from manas_ai.models import Document

async def test_custom_node():
    # Create node
    node = CustomQANode(name="test_node")
    
    # Test processing
    result = await node.process_question(
        "Test question?"
    )
    
    assert result["enhanced"]
    assert result["original_question"] == "Test question?"

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import numpy as np

async def test_custom_store():
    # Create store
    store = CustomVectorStore(dimension=3)
    
    # Add documents
    docs = [
        Document(content="Test 1"),
        Document(content="Test 2")
    ]
    count = await store.add_documents(docs)
    assert count == 2
    
    # Test search
    results = await store.similarity_search(
        "Test query",
        k=1
    )
    assert len(results) == 1
    assert isinstance(results[0][1], float)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import pytest
from manas_ai.models import ChatMessage

async def test_custom_provider():
    # Create provider
    provider = CustomProvider(api_key="test_key")
    
    # Test generation
    response = await provider.generate(
        "Test prompt",
        max_tokens=50
    )
    assert isinstance(response, str)
    
    # Test chat
    messages = [
        ChatMessage(role="user", content="Hello")
    ]
    response = await provider.chat(messages)
    assert isinstance(response, str)