diff --git a/prometheus/exceptions/__init__.py b/prometheus/exceptions/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/prometheus/exceptions/file_operation_exceptions.py b/prometheus/exceptions/file_operation_exceptions.py
new file mode 100644
index 00000000..e9b49bc1
--- /dev/null
+++ b/prometheus/exceptions/file_operation_exceptions.py
@@ -0,0 +1,6 @@
+class FileOperationException(Exception):
+    """
+    Base class for file operation exceptions.
+    """
+
+    pass
diff --git a/prometheus/lang_graph/nodes/context_extraction_node.py b/prometheus/lang_graph/nodes/context_extraction_node.py
new file mode 100644
index 00000000..bd1f7979
--- /dev/null
+++ b/prometheus/lang_graph/nodes/context_extraction_node.py
@@ -0,0 +1,130 @@
+import logging
+from typing import Sequence
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import SystemMessage
+from pydantic import BaseModel, Field
+
+from prometheus.exceptions.file_operation_exceptions import FileOperationException
+from prometheus.lang_graph.subgraphs.context_retrieval_state import ContextRetrievalState
+from prometheus.models.context import Context
+from prometheus.utils.file_utils import read_file_with_line_numbers
+
+SYS_PROMPT = """\
+You are a context summary agent that summarizes code context that is relevant to a given query from history messages.
+ Your goal is to extract, evaluate and summary code context that directly answers the query requirements.
+
+Your evaluation and summarization must consider two key aspects:
+1. Query Match: Which set of history messages directly address specific requirements mentioned in the query?
+2. Extended relevance: Which set of history messages provide essential information needed to understand the query topic?
+
+Follow these strict evaluation steps:
+1. First, identify specific requirements in the query
+2. Check which set of history messages directly addresses these requirements
+3. Check which parts of code context are relevant to the query
+4. Consider if they provides essential context by examining:
+   - Function dependencies
+   - Type definitions
+   - Configuration requirements
+   - Implementation details needed for completeness
+
+Query relevance guidelines - include only if:
+- It directly implements functionality mentioned in the query
+- It contains specific elements the query asks about
+- It's necessary to understand or implement query requirements
+- It provides critical information needed to answer the query
+
+CRITICAL RULE:
+- You don't have to select whole piece of code that you have seen, ONLY select the parts that are relevant to the query.
+- Each context MUST be SHORT and CONCISE, focusing ONLY on the lines that are relevant to the query.
+- Several context can be extracted from the same file, but each context must be concise and relevant to the query.
+- Do NOT include any irrelevant lines or comments that do not contribute to answering the query.
+- Do NOT include same context multiple times, even if it appears in different history messages.
+
+Remember: Your primary goal is to summarize context that directly helps answer the query requirements.
+
+Provide your analysis in a structured format matching the ContextExtractionStructuredOutput model.
+
+Example output:
+```json
+{
+    "context": [{
+        "reasoning": "1. Query requirement analysis:\n   - Query specifically asks about password validation\n   - Context provides implementation details for password validation\n2. Extended relevance:\n   - This function is essential for understanding how passwords are validated in the system",
+        "relative_path": "pychemia/code/fireball/fireball.py",
+        "start_line": 270,
+        "end_line": 293
+    } ......]
+}
+```
+
+Your task is to summarize the context from the provided history messages and return it in the specified format.
+"""
+
+
+class ContextOutput(BaseModel):
+    reasoning: str = Field(
+        description="Your step-by-step reasoning why the context is relevant to the query"
+    )
+    relative_path: str = Field(description="Relative path to the context file in the codebase")
+    start_line: int = Field(
+        description="Start line number of the context in the file, minimum is 1", ge=1
+    )
+    end_line: int = Field(
+        description="End line number of the context in the file, minimum is 1. "
+        "The Content in the end line is including",
+        ge=1,
+    )
+
+
+class ContextExtractionStructuredOutput(BaseModel):
+    context: Sequence[ContextOutput]
+
+
+class ContextExtractionNode:
+    def __init__(self, model: BaseChatModel, root_path: str):
+        structured_llm = model.with_structured_output(ContextExtractionStructuredOutput)
+        self.model = structured_llm
+        self.system_prompt = SystemMessage(SYS_PROMPT)
+        self.root_path = root_path
+        self._logger = logging.getLogger("prometheus.lang_graph.nodes.context_extraction_node")
+
+    def __call__(self, state: ContextRetrievalState):
+        self._logger.info("Starting context extraction process")
+        # Get Context List with existing context
+        final_context = state.get("context", [])
+        # Get chat history messages from the state
+        last_messages = state["context_provider_messages"]
+        # Summarize the context based on the last messages and system prompt
+        response = self.model.invoke([self.system_prompt] + last_messages)
+        self._logger.debug(f"Model response: {response}")
+        context_list = response.context
+        for context_ in context_list:
+            try:
+                content = read_file_with_line_numbers(
+                    relative_path=context_.relative_path,
+                    root_path=str(self.root_path),
+                    start_line=context_.start_line,
+                    end_line=context_.end_line,
+                )
+            except FileOperationException as e:
+                self._logger.error(e)
+                continue
+            if content:
+                final_context.append(
+                    Context(
+                        relative_path=context_.relative_path,
+                        start_line_number=context_.start_line,
+                        end_line_number=context_.end_line,
+                        content=content,
+                    )
+                )
+        # Filter out duplicate Context entries
+        seen = set()
+        unique_context = []
+        for ctx in final_context:
+            key = (ctx.relative_path, ctx.start_line_number, ctx.end_line_number)
+            if key not in seen:
+                seen.add(key)
+                unique_context.append(ctx)
+        self._logger.info(f"Context extraction complete, returning context {unique_context}")
+        return {"context": unique_context}
diff --git a/prometheus/lang_graph/nodes/context_provider_node.py b/prometheus/lang_graph/nodes/context_provider_node.py
index e7401f9d..30fd8345 100644
--- a/prometheus/lang_graph/nodes/context_provider_node.py
+++ b/prometheus/lang_graph/nodes/context_provider_node.py
@@ -34,7 +34,7 @@ class ContextProviderNode:
 
     SYS_PROMPT = """\
 You are a context gatherer that searches a Neo4j knowledge graph representation of a 
-codebase. Your role is to efficiently find relevant code and documentation 
+codebase. Your role is to understand the logic of the project and efficiently find relevant code and documentation 
 context based on user queries.
 
 Knowledge Graph Structure:
@@ -55,7 +55,7 @@ class ContextProviderNode:
    - Prioritize relative_path tools when exact file location is known
    - Fall back to basename tools for filename-only searches
    - Use AST node searches to find specific code structures
-   - Use preview_* or read_* tools with more than hundrend lines to get more context than class/function
+   - Use preview_* or read_* tools with more than hundred lines to get more context than class/function
    - If a search returns no results, try alternative approaches with broader scope
 
 2. Documentation/Text Search:
@@ -72,11 +72,8 @@ class ContextProviderNode:
 
 4. Critical Rules:
    - Do not repeat the same query!
-   - Do not select a whole file or directory as context, but rather specific code snippets.
-   - Each context should be a small, focused piece of code or documentation that directly addresses the query, which must be less than 100 lines!
-   - But several contexts snippets can be selected if they are relevant to the query.
 
-In your response, just provide a short summary with a few setences (3-4 setences) on what you have done.
+In your response, just provide a short summary with a few sentences (3-4 sentences) on what you have done.
 As your searched are automatically visible to the user, you do not need to repeat them. 
 
 The file tree of the codebase:
diff --git a/prometheus/lang_graph/nodes/context_refine_node.py b/prometheus/lang_graph/nodes/context_refine_node.py
index 648b55c2..05861bcf 100644
--- a/prometheus/lang_graph/nodes/context_refine_node.py
+++ b/prometheus/lang_graph/nodes/context_refine_node.py
@@ -69,9 +69,11 @@ class ContextRefineStructuredOutput(BaseModel):
 """
 
     def __init__(self, model: BaseChatModel, kg: KnowledgeGraph):
+        file_tree = kg.get_file_tree().replace("{", "{{").replace("}", "}}")
+        system_prompt = self.SYS_PROMPT.format(file_tree=file_tree)
         prompt = ChatPromptTemplate.from_messages(
             [
-                ("system", self.SYS_PROMPT.format(file_tree=kg.get_file_tree())),
+                ("system", system_prompt),
                 ("human", "{human_prompt}"),
             ]
         )
diff --git a/prometheus/lang_graph/nodes/context_selection_node.py b/prometheus/lang_graph/nodes/context_selection_node.py
deleted file mode 100644
index 6286c5d0..00000000
--- a/prometheus/lang_graph/nodes/context_selection_node.py
+++ /dev/null
@@ -1,112 +0,0 @@
-import logging
-
-from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.prompts import ChatPromptTemplate
-from pydantic import BaseModel, Field
-
-from prometheus.lang_graph.subgraphs.context_retrieval_state import ContextRetrievalState
-from prometheus.utils.lang_graph_util import extract_last_tool_messages
-from prometheus.utils.neo4j_util import neo4j_data_for_context_generator
-
-
-class ContextSelectionStructuredOutput(BaseModel):
-    reasoning: str = Field(
-        description="Your step-by-step reasoning why the context is relevant to the query"
-    )
-    relevant: bool = Field(description="If the context is relevant to the query")
-
-
-class ContextSelectionNode:
-    SYS_PROMPT = """\
-You are a context selection agent that evaluates if a piece of code context is relevant to a given query. Your goal is to determine if the context directly answers the query requirements.
-
-Your evaluation must consider two key aspects:
-1. Query Match: Does the context directly address specific requirements mentioned in the query?
-2. Extended relevance: Does this context provide essential information needed to understand the query topic?
-
-Follow these strict evaluation steps:
-1. First, identify specific requirements in the query
-2. Check if the context directly addresses these requirements
-3. Consider if it provides essential context by examining:
-   - Function dependencies
-   - Type definitions
-   - Configuration requirements
-   - Implementation details needed for completeness
-
-Query relevance guidelines - include only if:
-- It directly implements functionality mentioned in the query
-- It contains specific elements the query asks about
-- It's necessary to understand or implement query requirements
-- It provides critical information needed to answer the query
-
-Remember: Your primary goal is to determine if this specific piece of context directly helps answer the query requirements.
-
-Provide your analysis in a structured format matching the ContextSelectionStructuredOutput model.
-
-Example:
-
-Query: "How does the login endpoint validate passwords?"
-
-Context to evaluate:
-```python
-def validate_password(password: str, hash: str):
-    return bcrypt.checkpw(password.encode(), hash.encode())
-```
-
-Example output:
-```json
-{
-  "reasoning": "1. Query requirement analysis:
-   - Query specifically asks about password validation
-   - Needs implementation details of validation process
-   2. Context evaluation:
-   - Directly implements the password validation function
-   - Shows exactly how passwords are compared using bcrypt
-   - Provides essential implementation detail
-   3. Relevance confirmation:
-   - Directly answers how password validation works
-   - Shows specific security mechanism (bcrypt) used",
-  "relevant": true
-}
-```
-
-Your task is to analyze the context and provide a similar structured output with detailed reasoning and a relevance decision.
-""".replace("{", "{{").replace("}", "}}")
-
-    HUMAN_PROMPT = """\
-Query:
-{query}
-
-Found context:
-{context}
-
-Please classify if the found context is relevant to the query.
-"""
-
-    def __init__(self, model: BaseChatModel):
-        prompt = ChatPromptTemplate.from_messages(
-            [("system", self.SYS_PROMPT), ("human", "{human_prompt}")]
-        )
-        structured_llm = model.with_structured_output(ContextSelectionStructuredOutput)
-        self.model = prompt | structured_llm
-        self._logger = logging.getLogger("prometheus.lang_graph.nodes.context_selection_node")
-
-    def format_human_prompt(self, state: ContextRetrievalState, search_result: str) -> str:
-        context_info = self.HUMAN_PROMPT.format(query=state["query"], context=search_result)
-        return context_info
-
-    def __call__(self, state: ContextRetrievalState):
-        self._logger.info("Starting context selection process")
-        context_list = state.get("context", [])
-        for tool_message in extract_last_tool_messages(state["context_provider_messages"]):
-            for context in neo4j_data_for_context_generator(tool_message.artifact):
-                context_str = str(context)
-                human_prompt = self.format_human_prompt(state, context_str)
-                response = self.model.invoke({"human_prompt": human_prompt})
-                self._logger.debug(
-                    f"Is this search result {context_str} relevant?: {response.relevant}"
-                )
-                if response.relevant:
-                    context_list.append(context)
-        self._logger.info(f"Context selection complete, returning context {context_list}")
-        return {"context": context_list}
diff --git a/prometheus/lang_graph/nodes/final_patch_selection_node.py b/prometheus/lang_graph/nodes/final_patch_selection_node.py
index 09bd651d..5637a3c5 100644
--- a/prometheus/lang_graph/nodes/final_patch_selection_node.py
+++ b/prometheus/lang_graph/nodes/final_patch_selection_node.py
@@ -147,7 +147,9 @@ def __call__(self, state: Dict):
         human_prompt = self.format_human_message(state)
         for try_index in range(self.max_retries):
             response = self.model.invoke({"human_prompt": human_prompt})
-            self._logger.info(f"FinalPatchSelectionNode response at {try_index} try:\n{response}")
+            self._logger.info(
+                f"FinalPatchSelectionNode response at {try_index + 1} try:\n{response}"
+            )
 
             if 0 <= response.patch_index < len(state["edit_patches"]):
                 return {"final_patch": state["edit_patches"][response.patch_index]}
diff --git a/prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py b/prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py
index 8a4893ff..656c4adc 100644
--- a/prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py
@@ -182,7 +182,7 @@ def invoke(
             "issue_title": issue_title,
             "issue_body": issue_body,
             "issue_comments": issue_comments,
-            "max_refined_query_loop": 1,
+            "max_refined_query_loop": 3,
         }
 
         try:
diff --git a/prometheus/lang_graph/subgraphs/context_retrieval_subgraph.py b/prometheus/lang_graph/subgraphs/context_retrieval_subgraph.py
index 7f024827..a337588f 100644
--- a/prometheus/lang_graph/subgraphs/context_retrieval_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/context_retrieval_subgraph.py
@@ -7,10 +7,10 @@
 from langgraph.prebuilt import ToolNode, tools_condition
 
 from prometheus.graph.knowledge_graph import KnowledgeGraph
+from prometheus.lang_graph.nodes.context_extraction_node import ContextExtractionNode
 from prometheus.lang_graph.nodes.context_provider_node import ContextProviderNode
 from prometheus.lang_graph.nodes.context_query_message_node import ContextQueryMessageNode
 from prometheus.lang_graph.nodes.context_refine_node import ContextRefineNode
-from prometheus.lang_graph.nodes.context_selection_node import ContextSelectionNode
 from prometheus.lang_graph.nodes.reset_messages_node import ResetMessagesNode
 from prometheus.lang_graph.subgraphs.context_retrieval_state import ContextRetrievalState
 from prometheus.models.context import Context
@@ -69,8 +69,8 @@ def __init__(
             messages_key="context_provider_messages",
         )
 
-        # Step 4: Select relevant context snippets from the candidates
-        context_selection_node = ContextSelectionNode(model)
+        # Step 4: Extract the Context
+        context_extraction_node = ContextExtractionNode(model, str(kg.get_local_path()))
 
         # Step 5: Reset tool messages to prepare for the next iteration (if needed)
         reset_context_provider_messages_node = ResetMessagesNode("context_provider_messages")
@@ -85,7 +85,7 @@ def __init__(
         workflow.add_node("context_query_message_node", context_query_message_node)
         workflow.add_node("context_provider_node", context_provider_node)
         workflow.add_node("context_provider_tools", context_provider_tools)
-        workflow.add_node("context_selection_node", context_selection_node)
+        workflow.add_node("context_extraction_node", context_extraction_node)
         workflow.add_node(
             "reset_context_provider_messages_node", reset_context_provider_messages_node
         )
@@ -100,10 +100,10 @@ def __init__(
         workflow.add_conditional_edges(
             "context_provider_node",
             functools.partial(tools_condition, messages_key="context_provider_messages"),
-            {"tools": "context_provider_tools", END: "context_selection_node"},
+            {"tools": "context_provider_tools", END: "context_extraction_node"},
         )
         workflow.add_edge("context_provider_tools", "context_provider_node")
-        workflow.add_edge("context_selection_node", "reset_context_provider_messages_node")
+        workflow.add_edge("context_extraction_node", "reset_context_provider_messages_node")
         workflow.add_edge("reset_context_provider_messages_node", "context_refine_node")
 
         # If refined_query is non-empty, loop back to provider; else terminate
@@ -129,7 +129,7 @@ def invoke(
 
         Returns:
             Dict with a single key:
-                - "context" (Sequence[str]): A list of selected context snippets relevant to the query.
+                - "context" (Sequence[Context]): A list of selected context snippets relevant to the query.
         """
         config = {"recursion_limit": recursion_limit}
 
diff --git a/prometheus/lang_graph/subgraphs/issue_classification_subgraph.py b/prometheus/lang_graph/subgraphs/issue_classification_subgraph.py
index c6b2b73f..0844456d 100644
--- a/prometheus/lang_graph/subgraphs/issue_classification_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/issue_classification_subgraph.py
@@ -60,7 +60,7 @@ def invoke(
             "issue_title": issue_title,
             "issue_body": issue_body,
             "issue_comments": issue_comments,
-            "max_refined_query_loop": 1,
+            "max_refined_query_loop": 3,
         }
 
         output_state = self.subgraph.invoke(
diff --git a/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py b/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py
index 6346ea29..f7c4417a 100644
--- a/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py
@@ -123,7 +123,7 @@ def invoke(
             "issue_body": issue_body,
             "issue_comments": issue_comments,
             "number_of_candidate_patch": number_of_candidate_patch,
-            "max_refined_query_loop": 3,
+            "max_refined_query_loop": 5,
         }
 
         output_state = self.subgraph.invoke(input_state, config)
diff --git a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
index e151116a..2611402c 100644
--- a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
@@ -197,7 +197,7 @@ def invoke(
             "run_existing_test": run_existing_test,
             "reproduced_bug_file": reproduced_bug_file,
             "reproduced_bug_commands": reproduced_bug_commands,
-            "max_refined_query_loop": 3,
+            "max_refined_query_loop": 5,
         }
 
         output_state = self.subgraph.invoke(input_state, config)
diff --git a/prometheus/tools/graph_traversal.py b/prometheus/tools/graph_traversal.py
index 20fc9bc1..d2d6d790 100644
--- a/prometheus/tools/graph_traversal.py
+++ b/prometheus/tools/graph_traversal.py
@@ -1,10 +1,12 @@
 from pathlib import Path
+from typing import Any, Mapping, Sequence, Union
 
 from neo4j import GraphDatabase
 from pydantic import BaseModel, Field
 
 from prometheus.parser import tree_sitter_parser
 from prometheus.utils import neo4j_util
+from prometheus.utils.str_util import pre_append_line_numbers
 
 MAX_RESULT = 30
 
@@ -39,7 +41,7 @@ class FindFileNodeWithBasenameInput(BaseModel):
 
 def find_file_node_with_basename(
     basename: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""
       MATCH (f:FileNode {{ basename: '{basename}' }})
       RETURN f AS FileNode
@@ -64,7 +66,7 @@ class FindFileNodeWithRelativePathInput(BaseModel):
 
 def find_file_node_with_relative_path(
     relative_path: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
     MATCH (f:FileNode {{ relative_path: '{relative_path}' }})
     RETURN f AS FileNode
@@ -95,7 +97,7 @@ class FindASTNodeWithTextInFileWithBasenameInput(BaseModel):
 
 def find_ast_node_with_text_in_file_with_basename(
     text: str, basename: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
     MATCH (f:FileNode) -[:HAS_FILE*0..]-> (c:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*0..]-> (a:ASTNode)
     WHERE f.basename = '{basename}' AND a.text CONTAINS '{text}'
@@ -122,7 +124,7 @@ class FindASTNodeWithTextInFileWithRelativePathInput(BaseModel):
 
 def find_ast_node_with_text_in_file_with_relative_path(
     text: str, relative_path: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
         MATCH (f:FileNode) -[:HAS_FILE*0..]-> (c:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*0..]-> (a:ASTNode)
         WHERE f.relative_path = '{relative_path}' AND a.text CONTAINS '{text}'
@@ -147,7 +149,7 @@ class FindASTNodeWithTypeInFileWithBasenameInput(BaseModel):
 
 def find_ast_node_with_type_in_file_with_basename(
     type: str, basename: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
     MATCH (f:FileNode) -[:HAS_FILE*0..]-> (c:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*0..]-> (a:ASTNode)
     WHERE f.basename = '{basename}' AND a.type = '{type}'
@@ -172,7 +174,7 @@ class FindASTNodeWithTypeInFileWithRelativePathInput(BaseModel):
 
 def find_ast_node_with_type_in_file_with_relative_path(
     type: str, relative_path: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
         MATCH (f:FileNode) -[:HAS_FILE*0..]-> (c:FileNode) -[:HAS_AST]-> (:ASTNode) -[:PARENT_OF*0..]-> (a:ASTNode)
         WHERE f.relative_path = '{relative_path}' AND a.type = '{type}'
@@ -202,7 +204,7 @@ class FindTextNodeWithTextInput(BaseModel):
 
 def find_text_node_with_text(
     text: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
     MATCH (f:FileNode) -[:HAS_TEXT]-> (t:TextNode)
     WHERE t.text CONTAINS '{text}'
@@ -230,7 +232,7 @@ class FindTextNodeWithTextInFileInput(BaseModel):
 
 def find_text_node_with_text_in_file(
     text: str, basename: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
     MATCH (f:FileNode) -[:HAS_TEXT]-> (t:TextNode)
     WHERE f.basename = '{basename}' AND t.text CONTAINS '{text}'
@@ -253,7 +255,7 @@ class GetNextTextNodeWithNodeIdInput(BaseModel):
 
 def get_next_text_node_with_node_id(
     node_id: int, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     query = f"""\
     MATCH (f:FileNode) -[:HAS_TEXT]-> (a:TextNode {{ node_id: {node_id} }}) -[:NEXT_CHUNK]-> (b:TextNode)
     RETURN f as FileNode, b AS TextNode
@@ -282,7 +284,7 @@ class PreviewFileContentWithBasenameInput(BaseModel):
 
 def preview_file_content_with_basename(
     basename: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     source_code_query = f"""\
     MATCH (f:FileNode {{ basename: '{basename}' }}) -[:HAS_AST]-> (a:ASTNode)
     WITH f, apoc.text.split(a.text, '\\R') AS lines
@@ -304,8 +306,18 @@ def preview_file_content_with_basename(
   """
 
     if tree_sitter_parser.supports_file(Path(basename)):
-        return neo4j_util.run_neo4j_query(source_code_query, driver, max_token_per_result)
-    return neo4j_util.run_neo4j_query(text_query, driver, max_token_per_result)
+        data = neo4j_util.run_neo4j_query_without_formatting(source_code_query, driver)
+    else:
+        data = neo4j_util.run_neo4j_query_without_formatting(text_query, driver)
+    for result in data:
+        if isinstance(result["preview"], dict):
+            result["preview"]["text"] = pre_append_line_numbers(
+                result["preview"]["text"], result["preview"]["start_line"]
+            )
+            result["preview"]["end_line"] = (
+                result["preview"]["start_line"] + len(result["preview"]["text"].splitlines()) - 1
+            )
+    return neo4j_util.format_neo4j_data(data, max_token_per_result), data
 
 
 class PreviewFileContentWithRelativePathInput(BaseModel):
@@ -324,7 +336,7 @@ class PreviewFileContentWithRelativePathInput(BaseModel):
 
 def preview_file_content_with_relative_path(
     relative_path: str, driver: GraphDatabase.driver, max_token_per_result: int
-) -> str:
+) -> tuple[str, Sequence[Mapping[str, Any]]]:
     source_code_query = f"""\
       MATCH (f:FileNode {{ relative_path: '{relative_path}' }}) -[:HAS_AST]-> (a:ASTNode)
       WITH f, apoc.text.split(a.text, '\\R') AS lines
@@ -346,8 +358,18 @@ def preview_file_content_with_relative_path(
   """
 
     if tree_sitter_parser.supports_file(Path(relative_path)):
-        return neo4j_util.run_neo4j_query(source_code_query, driver, max_token_per_result)
-    return neo4j_util.run_neo4j_query(text_query, driver, max_token_per_result)
+        data = neo4j_util.run_neo4j_query_without_formatting(source_code_query, driver)
+    else:
+        data = neo4j_util.run_neo4j_query_without_formatting(text_query, driver)
+    for result in data:
+        if isinstance(result["preview"], dict):
+            result["preview"]["text"] = pre_append_line_numbers(
+                result["preview"]["text"], result["preview"]["start_line"]
+            )
+            result["preview"]["end_line"] = (
+                result["preview"]["start_line"] + len(result["preview"]["text"].splitlines()) - 1
+            )
+    return neo4j_util.format_neo4j_data(data, max_token_per_result), data
 
 
 class ReadCodeWithBasenameInput(BaseModel):
@@ -378,9 +400,9 @@ def read_code_with_basename(
     end_line: int,
     driver: GraphDatabase.driver,
     max_token_per_result: int,
-) -> str:
+) -> tuple[str, Union[Sequence[Mapping[str, Any]], None]]:
     if end_line < start_line:
-        return f"end_line {end_line} must be greater than start_line {start_line}"
+        return f"end_line {end_line} must be greater than start_line {start_line}", None
 
     source_code_query = f"""\
     MATCH (f:FileNode {{ basename: '{basename}' }}) -[:HAS_AST]-> (a:ASTNode)
@@ -394,8 +416,12 @@ def read_code_with_basename(
       }} AS SelectedLines
     ORDER BY f.node_id
   """
-
-    return neo4j_util.run_neo4j_query(source_code_query, driver, max_token_per_result)
+    data = neo4j_util.run_neo4j_query_without_formatting(source_code_query, driver)
+    for result in data:
+        result["SelectedLines"]["text"] = pre_append_line_numbers(
+            result["SelectedLines"]["text"], result["SelectedLines"]["start_line"]
+        )
+    return neo4j_util.format_neo4j_data(data, max_token_per_result), data
 
 
 class ReadCodeWithRelativePathInput(BaseModel):
@@ -427,9 +453,9 @@ def read_code_with_relative_path(
     end_line: int,
     driver: GraphDatabase.driver,
     max_token_per_result: int,
-) -> str:
+) -> tuple[str, Union[Sequence[Mapping[str, Any]], None]]:
     if end_line < start_line:
-        return f"end_line {end_line} must be greater than start_line {start_line}"
+        return f"end_line {end_line} must be greater than start_line {start_line}", None
 
     source_code_query = f"""\
         MATCH (f:FileNode {{ relative_path: '{relative_path}' }}) -[:HAS_AST]-> (a:ASTNode)
@@ -444,4 +470,10 @@ def read_code_with_relative_path(
         ORDER BY f.node_id
     """
 
-    return neo4j_util.run_neo4j_query(source_code_query, driver, max_token_per_result)
+    data = neo4j_util.run_neo4j_query_without_formatting(source_code_query, driver)
+    for result in data:
+        result["SelectedLines"]["text"] = pre_append_line_numbers(
+            result["SelectedLines"]["text"], result["SelectedLines"]["start_line"]
+        )
+
+    return neo4j_util.format_neo4j_data(data, max_token_per_result), data
diff --git a/prometheus/utils/file_utils.py b/prometheus/utils/file_utils.py
new file mode 100644
index 00000000..fb3252af
--- /dev/null
+++ b/prometheus/utils/file_utils.py
@@ -0,0 +1,35 @@
+import os
+from pathlib import Path
+
+from prometheus.exceptions.file_operation_exceptions import FileOperationException
+
+
+def read_file_with_line_numbers(
+    relative_path: str, root_path: str, start_line: int, end_line: int
+) -> str:
+    if os.path.isabs(relative_path):
+        raise FileOperationException(
+            f"relative_path: {relative_path} is a absolute path, not relative path."
+        )
+
+    file_path = Path(os.path.join(root_path, relative_path))
+    if not file_path.exists():
+        raise FileOperationException(f"The file {relative_path} does not exist.")
+
+    if not file_path.is_file():
+        raise FileOperationException(f"The path {relative_path} is not a file.")
+
+    if end_line < start_line:
+        raise FileOperationException(
+            f"The end line number {end_line} must be greater than "
+            f"the start line number {start_line}."
+        )
+
+    zero_based_start_line = start_line - 1
+    # The content in the end line is included
+    zero_based_end_line = end_line
+
+    with file_path.open() as f:
+        lines = f.readlines()
+
+    return "".join(lines[zero_based_start_line:zero_based_end_line])
diff --git a/prometheus/utils/neo4j_util.py b/prometheus/utils/neo4j_util.py
index 5ef4263b..b594c5a9 100644
--- a/prometheus/utils/neo4j_util.py
+++ b/prometheus/utils/neo4j_util.py
@@ -1,18 +1,17 @@
-from typing import Any, Iterator, Mapping, Optional, Sequence, Tuple
+from typing import Any, Mapping, Sequence, Tuple
 
 import neo4j
 
-from prometheus.models.context import Context
 from prometheus.utils.str_util import truncate_text
 
-EMPTY_DATA_MESSAGE = "Your query returned empty result, please try a different query."
+EMPTY_DATA_MESSAGE = "Your query returned empty result, please try a different query!"
 
 
 def format_neo4j_data(data: Sequence[Mapping[str, Any]], max_token_per_result: int) -> str:
     """Format a Neo4j result into a string.
 
     Args:
-      result: The result from a Neo4j query.
+      data: The result from a Neo4j query.
       max_token_per_result: Maximum number of tokens per result.
 
     Returns:
@@ -30,35 +29,6 @@ def format_neo4j_data(data: Sequence[Mapping[str, Any]], max_token_per_result: i
     return truncate_text(output.strip(), max_token_per_result)
 
 
-def neo4j_data_for_context_generator(
-    data: Optional[Sequence[Mapping[str, Any]]],
-) -> Iterator[Context]:
-    if data is None:
-        return
-
-    for search_result in data:
-        search_result_keys = search_result.keys()
-        # Skip if the result has no keys or only contains the "FileNode" key
-        if len(search_result_keys) == 1:
-            continue
-
-        context = Context(
-            relative_path=search_result["FileNode"]["relative_path"],
-            content=search_result.get("ASTNode", {}).get("text")
-            or search_result.get("TextNode", {}).get("text")
-            or search_result.get("preview", {}).get("text")
-            or search_result.get("SelectedLines", {}).get("text"),
-            start_line_number=search_result.get("ASTNode", {}).get("start_line")
-            or search_result.get("SelectedLines", {}).get("start_line")
-            or search_result.get("preview", {}).get("start_line"),
-            end_line_number=search_result.get("ASTNode", {}).get("end_line")
-            or search_result.get("SelectedLines", {}).get("end_line")
-            or search_result.get("preview", {}).get("end_line"),
-        )
-
-        yield context
-
-
 def run_neo4j_query(
     query: str, driver: neo4j.GraphDatabase.driver, max_token_per_result: int
 ) -> Tuple[str, Sequence[Mapping[str, Any]]]:
@@ -80,3 +50,25 @@ def query_transaction(tx):
 
     with driver.session() as session:
         return session.execute_read(query_transaction)
+
+
+def run_neo4j_query_without_formatting(
+    query: str, driver: neo4j.GraphDatabase.driver
+) -> Sequence[Mapping[str, Any]]:
+    """Run a read-only Neo4j query and return the result.
+
+    Args:
+      query: The query to run.
+      driver: The Neo4j driver to use.
+
+    Returns:
+      result
+    """
+
+    def query_transaction(tx):
+        result = tx.run(query)
+        data = result.data()
+        return data
+
+    with driver.session() as session:
+        return session.execute_read(query_transaction)