Building Custom AI Coding Assistants with LangChain

General-purpose AI coding assistants like GitHub Copilot and ChatGPT are powerful, but they lack knowledge of your specific codebase, architectural decisions, and internal APIs. Building a custom AI coding assistant tailored to your project can dramatically improve developer productivity while maintaining consistency with your team's conventions.

In this comprehensive guide, we'll explore how to build a project-specific AI coding assistant using LangChain, implementing Retrieval-Augmented Generation (RAG) for codebase understanding, setting up vector databases for semantic code search, and creating intelligent agents that can reason about and execute code.

Why Build Custom AI Coding Assistants?

While off-the-shelf AI tools are impressive, they have fundamental limitations for project-specific work:

• No knowledge of your codebase: They cannot reference your custom utilities, internal APIs, or architectural patterns
• Outdated information: Training data has a cutoff date, missing your latest code changes
• Generic suggestions: Responses follow general best practices rather than your team's conventions
• Context limitations: Large codebases exceed context window limits

A custom assistant built with LangChain solves these problems by:

LangChain Fundamentals for Coding Assistants

LangChain is an open-source framework designed for building applications powered by language models. It provides essential abstractions that make building AI assistants straightforward.

Core LangChain Concepts

Before diving into implementation, let's understand the key components:

• LLMs and Chat Models: Interfaces to language models (OpenAI, Anthropic, local models)
• Prompts: Templates for structuring inputs to the model
• Chains: Sequences of calls combining prompts, LLMs, and processing
• Agents: Dynamic chains that decide which tools to use based on input
• Retrievers: Components that fetch relevant documents from vector stores
• Memory: Mechanisms for maintaining conversation context

Setting Up Your Environment

First, install the required packages:

# Install LangChain and dependencies
pip install langchain langchain-openai langchain-community
pip install chromadb  # Vector database
pip install tiktoken  # Token counting
pip install python-dotenv  # Environment variables

# For TypeScript/JavaScript projects
npm install langchain @langchain/openai @langchain/community
npm install chromadb

Create a basic configuration:

# .env file
OPENAI_API_KEY=your-openai-api-key
ANTHROPIC_API_KEY=your-anthropic-api-key  # Optional

# Python: config.py
import os
from dotenv import load_dotenv

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
EMBEDDING_MODEL = "text-embedding-3-small"
CHAT_MODEL = "gpt-4-turbo-preview"
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 200

Setting Up Vector Databases for Code Search

Vector databases store code as high-dimensional embeddings, enabling semantic search. When you ask "how does authentication work?", the system finds relevant code even if it doesn't contain those exact words.

Choosing a Vector Database

Popular options for code search include:

• Chroma: Lightweight, embedded, perfect for development and small projects
• Pinecone: Managed service, highly scalable, production-ready
• Weaviate: Supports hybrid search (vector + keyword), self-hosted option
• pgvector: PostgreSQL extension, integrates with existing databases

Implementing Code Indexing with Chroma

# code_indexer.py
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
import os
from pathlib import Path

class CodebaseIndexer:
    def __init__(self, project_path: str, persist_directory: str = "./chroma_db"):
        self.project_path = project_path
        self.persist_directory = persist_directory
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

        # Code-aware text splitter
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=[
                "\nclass ",      # Class definitions
                "\ndef ",        # Function definitions
                "\n\n",          # Paragraph breaks
                "\n",            # Line breaks
                " ",             # Word breaks
            ]
        )

    def get_code_files(self) -> list:
        """Get all code files from the project."""
        extensions = [
            "*.py", "*.js", "*.ts", "*.jsx", "*.tsx",
            "*.java", "*.go", "*.rs", "*.cpp", "*.c",
            "*.md", "*.json", "*.yaml", "*.yml"
        ]

        files = []
        for ext in extensions:
            files.extend(Path(self.project_path).rglob(ext))

        # Filter out common non-essential directories
        excluded_dirs = {"node_modules", ".git", "__pycache__", "venv", "dist", "build"}
        return [f for f in files if not any(d in f.parts for d in excluded_dirs)]

    def load_documents(self):
        """Load and split documents from the codebase."""
        documents = []

        for file_path in self.get_code_files():
            try:
                loader = TextLoader(str(file_path), encoding="utf-8")
                docs = loader.load()

                # Add metadata for better retrieval
                for doc in docs:
                    doc.metadata["source"] = str(file_path)
                    doc.metadata["file_type"] = file_path.suffix
                    doc.metadata["relative_path"] = str(file_path.relative_to(self.project_path))

                documents.extend(docs)
            except Exception as e:
                print(f"Error loading {file_path}: {e}")

        # Split documents into chunks
        return self.splitter.split_documents(documents)

    def create_vector_store(self):
        """Create and persist the vector store."""
        documents = self.load_documents()
        print(f"Indexing {len(documents)} document chunks...")

        vectorstore = Chroma.from_documents(
            documents=documents,
            embedding=self.embeddings,
            persist_directory=self.persist_directory
        )

        print(f"Vector store created with {vectorstore._collection.count()} vectors")
        return vectorstore

    def load_vector_store(self):
        """Load an existing vector store."""
        return Chroma(
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )

# Usage
if __name__ == "__main__":
    indexer = CodebaseIndexer(
        project_path="./my-project",
        persist_directory="./my-project-vectors"
    )
    vectorstore = indexer.create_vector_store()

Optimizing Embeddings for Code

Code has different semantic properties than natural language. Consider these optimizations:

# Enhanced code splitter with AST awareness
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter

class LanguageAwareCodeSplitter:
    """Split code using language-specific syntax awareness."""

    LANGUAGE_MAP = {
        ".py": Language.PYTHON,
        ".js": Language.JS,
        ".ts": Language.TS,
        ".java": Language.JAVA,
        ".go": Language.GO,
        ".rs": Language.RUST,
    }

    def get_splitter(self, file_extension: str):
        """Get appropriate splitter for the file type."""
        language = self.LANGUAGE_MAP.get(file_extension)

        if language:
            return RecursiveCharacterTextSplitter.from_language(
                language=language,
                chunk_size=1000,
                chunk_overlap=200
            )

        # Fallback for unknown languages
        return RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )

Implementing RAG for Codebase Understanding

RAG combines retrieval with generation, allowing your assistant to answer questions using actual code from your repository.

Basic RAG Chain

# rag_assistant.py
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

class CodebaseAssistant:
    def __init__(self, vectorstore, model_name: str = "gpt-4-turbo-preview"):
        self.vectorstore = vectorstore
        self.llm = ChatOpenAI(
            model=model_name,
            temperature=0.1  # Lower temperature for more consistent code
        )

        self.prompt_template = PromptTemplate(
            input_variables=["context", "question"],
            template="""You are an expert coding assistant with deep knowledge of this codebase.
Use the following code snippets to answer the question. If you cannot find the answer in the
provided context, say so clearly and provide general guidance.

IMPORTANT GUIDELINES:
- Reference specific files and line numbers when possible
- Follow the coding patterns and conventions shown in the context
- Suggest improvements that align with the existing architecture
- If suggesting new code, make it consistent with the project style

Context from codebase:
{context}

Question: {question}

Answer:"""
        )

        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",  # Use "map_reduce" for very large contexts
            retriever=self.vectorstore.as_retriever(
                search_type="similarity",
                search_kwargs={"k": 5}  # Retrieve top 5 relevant chunks
            ),
            chain_type_kwargs={"prompt": self.prompt_template},
            return_source_documents=True
        )

    def ask(self, question: str) -> dict:
        """Ask a question about the codebase."""
        result = self.qa_chain.invoke({"query": question})

        return {
            "answer": result["result"],
            "sources": [
                {
                    "file": doc.metadata.get("relative_path", "Unknown"),
                    "content_preview": doc.page_content[:200] + "..."
                }
                for doc in result["source_documents"]
            ]
        }

# Usage
from code_indexer import CodebaseIndexer

indexer = CodebaseIndexer("./my-project")
vectorstore = indexer.load_vector_store()
assistant = CodebaseAssistant(vectorstore)

response = assistant.ask("How does the authentication middleware work?")
print(response["answer"])
print("\nSources:")
for source in response["sources"]:
    print(f"  - {source['file']}")

Advanced RAG with Hybrid Search

Combine vector similarity with keyword matching for better results:

# hybrid_retriever.py
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import Chroma

class HybridCodeRetriever:
    """Combine semantic search with keyword search."""

    def __init__(self, documents, vectorstore: Chroma):
        # BM25 for keyword matching (good for function names, variables)
        self.bm25_retriever = BM25Retriever.from_documents(documents)
        self.bm25_retriever.k = 3

        # Vector retriever for semantic search
        self.vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

        # Ensemble combines both with weights
        self.ensemble_retriever = EnsembleRetriever(
            retrievers=[self.bm25_retriever, self.vector_retriever],
            weights=[0.4, 0.6]  # 40% keyword, 60% semantic
        )

    def get_relevant_documents(self, query: str):
        """Retrieve documents using hybrid search."""
        return self.ensemble_retriever.get_relevant_documents(query)

Creating Multi-Step Reasoning Chains

Complex coding tasks often require multiple steps. LangChain's chains allow you to decompose problems intelligently.

Code Analysis Chain

# analysis_chain.py
from langchain.chains import SequentialChain, LLMChain
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate

class CodeAnalysisChain:
    """Multi-step chain for comprehensive code analysis."""

    def __init__(self, vectorstore):
        self.llm = ChatOpenAI(model="gpt-4-turbo-preview", temperature=0)
        self.vectorstore = vectorstore

    def create_analysis_chain(self):
        # Step 1: Retrieve relevant code
        retrieval_prompt = PromptTemplate(
            input_variables=["query"],
            template="Find code related to: {query}\nReturn file paths and relevant functions."
        )

        # Step 2: Understand the code structure
        structure_prompt = PromptTemplate(
            input_variables=["code_context", "query"],
            template="""Analyze this code structure:
{code_context}

For the query: {query}

Identify:
1. Main components involved
2. Data flow between components
3. External dependencies
4. Potential issues or improvements"""
        )

        # Step 3: Generate recommendations
        recommendation_prompt = PromptTemplate(
            input_variables=["structure_analysis", "query"],
            template="""Based on this analysis:
{structure_analysis}

Original query: {query}

Provide:
1. Specific code changes with file locations
2. Step-by-step implementation guide
3. Potential edge cases to consider
4. Testing recommendations"""
        )

        structure_chain = LLMChain(
            llm=self.llm,
            prompt=structure_prompt,
            output_key="structure_analysis"
        )

        recommendation_chain = LLMChain(
            llm=self.llm,
            prompt=recommendation_prompt,
            output_key="recommendations"
        )

        return SequentialChain(
            chains=[structure_chain, recommendation_chain],
            input_variables=["code_context", "query"],
            output_variables=["structure_analysis", "recommendations"]
        )

    def analyze(self, query: str) -> dict:
        # First, retrieve relevant code
        docs = self.vectorstore.similarity_search(query, k=8)
        code_context = "\n\n".join([
            f"File: {doc.metadata.get('relative_path')}\n{doc.page_content}"
            for doc in docs
        ])

        chain = self.create_analysis_chain()
        result = chain.invoke({
            "code_context": code_context,
            "query": query
        })

        return result

Building Agents That Can Execute Code

Agents are the most powerful LangChain feature - they dynamically decide which tools to use based on the task.

Creating a Code-Executing Agent

# code_agent.py
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_openai import ChatOpenAI
from langchain.tools import Tool, StructuredTool
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from pydantic import BaseModel, Field
import subprocess
import tempfile
import os

class CodeExecutionInput(BaseModel):
    """Input for code execution tool."""
    code: str = Field(description="Python code to execute")
    timeout: int = Field(default=30, description="Execution timeout in seconds")

class FileSearchInput(BaseModel):
    """Input for file search tool."""
    query: str = Field(description="Search query for finding files")

class CodeAssistantAgent:
    def __init__(self, vectorstore, project_path: str):
        self.vectorstore = vectorstore
        self.project_path = project_path
        self.llm = ChatOpenAI(model="gpt-4-turbo-preview", temperature=0)

        self.tools = self._create_tools()
        self.agent = self._create_agent()

    def _create_tools(self) -> list:
        """Create tools available to the agent."""

        def search_codebase(query: str) -> str:
            """Search the codebase for relevant code."""
            docs = self.vectorstore.similarity_search(query, k=5)
            results = []
            for doc in docs:
                results.append(f"File: {doc.metadata.get('relative_path')}\n{doc.page_content}\n---")
            return "\n".join(results)

        def execute_python(code: str, timeout: int = 30) -> str:
            """Execute Python code safely in a sandbox."""
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
                f.write(code)
                f.flush()

                try:
                    result = subprocess.run(
                        ['python', f.name],
                        capture_output=True,
                        text=True,
                        timeout=timeout,
                        cwd=self.project_path
                    )
                    output = result.stdout if result.stdout else result.stderr
                    return f"Output:\n{output}" if output else "Code executed successfully (no output)"
                except subprocess.TimeoutExpired:
                    return "Error: Code execution timed out"
                except Exception as e:
                    return f"Error: {str(e)}"
                finally:
                    os.unlink(f.name)

        def run_tests(test_path: str = "") -> str:
            """Run project tests."""
            cmd = ['pytest', '-v']
            if test_path:
                cmd.append(test_path)

            try:
                result = subprocess.run(
                    cmd,
                    capture_output=True,
                    text=True,
                    timeout=120,
                    cwd=self.project_path
                )
                return result.stdout + result.stderr
            except Exception as e:
                return f"Error running tests: {str(e)}"

        def read_file(file_path: str) -> str:
            """Read a file from the project."""
            full_path = os.path.join(self.project_path, file_path)
            try:
                with open(full_path, 'r') as f:
                    return f.read()
            except Exception as e:
                return f"Error reading file: {str(e)}"

        return [
            Tool(
                name="search_codebase",
                func=search_codebase,
                description="Search the codebase for relevant code snippets. Use for finding existing implementations, patterns, or understanding how features work."
            ),
            StructuredTool.from_function(
                func=execute_python,
                name="execute_python",
                description="Execute Python code. Use for testing code snippets, running scripts, or verifying implementations.",
                args_schema=CodeExecutionInput
            ),
            Tool(
                name="run_tests",
                func=run_tests,
                description="Run pytest tests. Optionally specify a test file path. Use to verify code changes don't break existing functionality."
            ),
            Tool(
                name="read_file",
                func=read_file,
                description="Read a specific file from the project. Use when you need the complete content of a file."
            )
        ]

    def _create_agent(self):
        """Create the agent with tools."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """You are an expert coding assistant with access to the project codebase.
You can search code, execute Python, run tests, and read files.

Guidelines:
1. Always search the codebase first to understand existing patterns
2. Before suggesting changes, verify they don't break tests
3. Follow the project's existing coding conventions
4. Provide complete, working code solutions
5. Explain your reasoning step by step"""),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad")
        ])

        agent = create_openai_tools_agent(self.llm, self.tools, prompt)

        return AgentExecutor(
            agent=agent,
            tools=self.tools,
            verbose=True,
            max_iterations=10,
            handle_parsing_errors=True
        )

    def run(self, task: str) -> str:
        """Execute a task using the agent."""
        return self.agent.invoke({"input": task})

# Usage
agent = CodeAssistantAgent(vectorstore, "./my-project")
result = agent.run("""
Add input validation to the user registration function.
Search for the existing implementation, add validation for email and password,
then run the tests to make sure everything works.
""")

Integrating with Development Tools

VS Code Extension Integration

Create a simple server that VS Code can communicate with:

# assistant_server.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from code_indexer import CodebaseIndexer
from rag_assistant import CodebaseAssistant
import uvicorn

app = FastAPI(title="Code Assistant API")

# Initialize on startup
indexer = None
assistant = None

class QueryRequest(BaseModel):
    question: str
    project_path: str = "./my-project"

class IndexRequest(BaseModel):
    project_path: str

@app.post("/index")
async def index_project(request: IndexRequest):
    """Index or re-index a project."""
    global indexer, assistant

    indexer = CodebaseIndexer(request.project_path)
    vectorstore = indexer.create_vector_store()
    assistant = CodebaseAssistant(vectorstore)

    return {"status": "indexed", "project": request.project_path}

@app.post("/ask")
async def ask_question(request: QueryRequest):
    """Ask a question about the codebase."""
    if not assistant:
        raise HTTPException(status_code=400, detail="No project indexed")

    response = assistant.ask(request.question)
    return response

@app.get("/health")
async def health_check():
    return {"status": "healthy", "indexed": assistant is not None}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

Git Hook Integration for Auto-Indexing

#!/bin/bash
# .git/hooks/post-commit

echo "Updating code index..."
python -c "
from code_indexer import CodebaseIndexer
indexer = CodebaseIndexer('.', './chroma_db')
indexer.create_vector_store()
print('Index updated successfully')
"

Cost Optimization Strategies

Running custom AI assistants can be expensive. Here are strategies to reduce costs:

# cost_aware_assistant.py
from langchain_openai import ChatOpenAI
import tiktoken

class CostAwareAssistant:
    def __init__(self, vectorstore, budget_tokens: int = 4000):
        self.vectorstore = vectorstore
        self.budget_tokens = budget_tokens
        self.encoder = tiktoken.encoding_for_model("gpt-4")

        # Use cheaper model for simple queries
        self.simple_llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
        self.complex_llm = ChatOpenAI(model="gpt-4-turbo-preview", temperature=0)

    def count_tokens(self, text: str) -> int:
        return len(self.encoder.encode(text))

    def get_context_within_budget(self, query: str, docs) -> str:
        """Build context that fits within token budget."""
        context_parts = []
        total_tokens = self.count_tokens(query) + 500  # Reserve for prompt

        for doc in docs:
            doc_tokens = self.count_tokens(doc.page_content)
            if total_tokens + doc_tokens <= self.budget_tokens:
                context_parts.append(doc.page_content)
                total_tokens += doc_tokens
            else:
                break

        return "\n\n".join(context_parts)

    def determine_complexity(self, query: str) -> str:
        """Determine if query needs complex model."""
        complex_keywords = ["refactor", "architecture", "design", "optimize", "debug complex"]
        return "complex" if any(kw in query.lower() for kw in complex_keywords) else "simple"

    def ask(self, query: str) -> str:
        docs = self.vectorstore.similarity_search(query, k=10)
        context = self.get_context_within_budget(query, docs)

        llm = self.complex_llm if self.determine_complexity(query) == "complex" else self.simple_llm

        # Continue with RAG chain using selected model...

Measuring Custom Assistant Effectiveness

Track these metrics to compare your custom assistant against general-purpose tools:

# metrics_tracker.py
import time
from dataclasses import dataclass
from typing import Optional
import json

@dataclass
class QueryMetrics:
    query: str
    response_time: float
    tokens_used: int
    sources_retrieved: int
    user_rating: Optional[int] = None  # 1-5 scale
    was_helpful: Optional[bool] = None

class MetricsTracker:
    def __init__(self, log_file: str = "assistant_metrics.jsonl"):
        self.log_file = log_file
        self.metrics = []

    def track_query(self, query: str, response_time: float, tokens: int, sources: int):
        metric = QueryMetrics(
            query=query,
            response_time=response_time,
            tokens_used=tokens,
            sources_retrieved=sources
        )
        self.metrics.append(metric)

        with open(self.log_file, 'a') as f:
            f.write(json.dumps(metric.__dict__) + '\n')

    def add_feedback(self, query_index: int, rating: int, helpful: bool):
        if 0 <= query_index < len(self.metrics):
            self.metrics[query_index].user_rating = rating
            self.metrics[query_index].was_helpful = helpful

    def get_summary(self) -> dict:
        if not self.metrics:
            return {}

        return {
            "total_queries": len(self.metrics),
            "avg_response_time": sum(m.response_time for m in self.metrics) / len(self.metrics),
            "avg_tokens": sum(m.tokens_used for m in self.metrics) / len(self.metrics),
            "avg_sources": sum(m.sources_retrieved for m in self.metrics) / len(self.metrics),
            "helpfulness_rate": sum(1 for m in self.metrics if m.was_helpful) / len([m for m in self.metrics if m.was_helpful is not None]) if any(m.was_helpful is not None for m in self.metrics) else None
        }

Key Takeaways

Conclusion

Building a custom AI coding assistant with LangChain transforms how your team interacts with your codebase. Instead of generic suggestions, developers get context-aware answers grounded in your actual code, following your conventions and understanding your architecture.

Start simple with a basic RAG setup, then gradually add capabilities like multi-step reasoning chains and code-executing agents. The investment in building custom tooling pays dividends in developer productivity and code consistency.

For more on managing AI context effectively, see our article on Context Window Limitations. To explore AI-assisted testing workflows that complement your custom assistant, check out Automated Testing with AI.